Docstrings and type annotations (#51)

Author: kgaughan

pyproject.toml

@@ -108,7 +108,7 @@ select = [
   "PLE",
   "PLR",
   "PLW",
-  # "PT",
+  "PT",
   "PYI",
   "Q",
   "RET",

src/uwhoisd/__init__.py

@@ -1,6 +1,4 @@
-"""
-A 'universal' WHOIS proxy server.
-"""
+"""A 'universal' WHOIS proxy server."""
 
 import asyncio
 import configparser
@@ -11,20 +9,19 @@
 import re
 import socket
 import sys
+import typing as t
 
 from . import caching, client, server, utils
 
 USAGE = "Usage: %s <config>"
 
 PORT = socket.getservbyname("whois", "tcp")
 
-logger = logging.getLogger("uwhoisd")
+logger = logging.getLogger(__name__)
 
 
 class UWhois:
-    """
-    Universal WHOIS proxy.
-    """
+    """Universal WHOIS proxy."""
 
     __slots__ = (
         "conservative",
@@ -36,22 +33,21 @@ class UWhois:
         "suffix",
     )
 
-    def __init__(self):
-        """
-        Initialise the proxy.
-        """
+    def __init__(self) -> None:
         super().__init__()
-        self.suffix = None
-        self.overrides = {}
-        self.prefixes = {}
-        self.recursion_patterns = {}
-        self.registry_whois = False
-        self.page_feed = True
-        self.conservative = ()
-
-    def read_config(self, parser):
-        """
-        Read the configuration for this object from a config file.
+        self.suffix: t.Optional[str] = None
+        self.overrides: dict[str, str] = {}
+        self.prefixes: dict[str, str] = {}
+        self.recursion_patterns: dict[str, re.Pattern] = {}
+        self.registry_whois: bool = False
+        self.page_feed: bool = True
+        self.conservative: t.Sequence[str] = ()
+
+    def read_config(self, parser: utils.ConfigParser) -> None:
+        """Read the configuration for this object from a config file.
+
+        Args:
+            parser: The config parser to read from.
         """
         self.registry_whois = parser.get_bool("uwhoisd", "registry_whois")
         self.page_feed = parser.get_bool("uwhoisd", "page_feed")
@@ -64,34 +60,53 @@ def read_config(self, parser):
         for zone, pattern in parser.items("recursion_patterns"):
             self.recursion_patterns[zone] = re.compile(utils.decode_value(pattern), re.IGNORECASE)
 
-    def get_whois_server(self, zone):
-        """
-        Get the WHOIS server for the given zone.
+    def get_whois_server(self, zone: str) -> tuple[str, int]:
+        """Get the WHOIS server for the given zone.
+
+        Args:
+            zone: The zone to get the WHOIS server for.
+
+        Returns:
+            A tuple of the WHOIS server and port.
         """
         server = self.overrides.get(zone, f"{zone}.{self.suffix}")
         if ":" in server:
             server, port = server.split(":", 1)
-            port = int(port)
-        else:
-            port = PORT
-        return server, port
+            return server, int(port)
+        return server, PORT
 
-    def get_registrar_whois_server(self, zone, response):
-        """
-        Extract the registrar's WHOIS server from the registry response.
+    def get_registrar_whois_server(self, zone: str, response: str) -> t.Optional[str]:
+        """Extract the registrar's WHOIS server from the registry response.
+
+        Args:
+            zone: The zone being queried.
+            response: The response from the registry's WHOIS server.
+
+        Returns:
+            The registrar's WHOIS server, or None if not found.
         """
         matches = self.recursion_patterns[zone].search(response)
         return None if matches is None else matches.group("server")
 
-    def get_prefix(self, zone):
-        """
-        Get the prefix required when querying the servers for the given zone.
+    def get_prefix(self, zone: str) -> str:
+        """Get the prefix required when querying the servers for the given zone.
+
+        Args:
+            zone: The zone to get the prefix for.
+
+        Returns:
+            The prefix string.
         """
         return self.prefixes.get(zone, "")
 
     async def whois(self, query: str) -> str:
-        """
-        Query the appropriate WHOIS server.
+        """Query the appropriate WHOIS server.
+
+        Args:
+            query: The WHOIS query.
+
+        Returns:
+            The WHOIS response.
         """
         # Figure out the zone whose WHOIS server we're meant to be querying.
         for zone in self.conservative:
@@ -107,23 +122,21 @@ async def whois(self, query: str) -> str:
 
         # Thin registry? Query the registrar's WHOIS server.
         if zone in self.recursion_patterns:
-            server = self.get_registrar_whois_server(zone, response)
-            if server is not None:
+            registrar_server = self.get_registrar_whois_server(zone, response)
+            if registrar_server is not None:
                 if not self.registry_whois:
                     response = ""
                 elif self.page_feed:
                     # A form feed character so it's possible to find the split.
                     response += "\f"
-                logger.info("Recursive query to %s about %s", server, query)
-                response += await client.query_whois(server, port, query)
+                logger.info("Recursive query to %s about %s", registrar_server, query)
+                response += await client.query_whois(registrar_server, port, query)
 
         return response
 
 
-def main():
-    """
-    Execute the daemon.
-    """
+def main() -> int:
+    """Execute the daemon."""
     if len(sys.argv) != 2:
         print(USAGE % os.path.basename(sys.argv[0]), file=sys.stderr)
         return 1

src/uwhoisd/caching.py

@@ -1,24 +1,49 @@
-"""
-Caching support.
-"""
+"""Caching support."""
 
 import collections
 from importlib import metadata
 import logging
 import time
+import typing as t
 
-logger = logging.getLogger("uwhoisd")
+logger = logging.getLogger(__name__)
+
+
+class Cache(t.Protocol):
+    """A WHOIS cache protocol."""
+
+    def get(self, key: str) -> t.Optional[str]:
+        """Retrieve a value from the cache.
+
+        Args:
+            key: The cache key to look up.
+
+        Returns:
+            The cached value, or `None` if not found.
+        """
+
+    def set(self, key: str, value: str) -> None:
+        """Store a value in the cache.
+
+        Args:
+            key: The cache key to store the value under.
+            value: The value to store.
+        """
 
 
 class UnknownCacheError(Exception):
-    """
-    The supplied cache type name cannot be found.
-    """
+    """The supplied cache type name cannot be found."""
 
 
-def get_cache(cfg):
-    """
-    Attempt to load the configured cache.
+def get_cache(cfg: dict[str, str]) -> t.Optional[Cache]:
+    """Attempt to load the configured cache.
+
+    Args:
+        cfg: The cache configuration. Must contain a "type" key giving the
+            cache type name. Any other keys are passed as parameters to the
+            cache constructor.
+
+    Returns: The cache object, or `None` if caching is disabled.
     """
     cache_name = cfg.pop("type", "null")
     if cache_name == "null":
@@ -33,14 +58,23 @@ def get_cache(cfg):
     raise UnknownCacheError(cache_name)
 
 
-def wrap_whois(cache, whois_func):
-    """
-    Wrap a WHOIS query function with a cache.
+def wrap_whois(
+    cache: t.Optional[Cache],
+    whois_func: t.Callable[[str], t.Awaitable[str]],
+) -> t.Callable[[str], t.Awaitable[str]]:
+    """Wrap a WHOIS query function with a cache.
+
+    Args:
+        cache: The cache to use, or `None` to disable caching.
+        whois_func: The WHOIS query function to wrap.
+
+    Returns:
+        The wrapped WHOIS query function.
     """
     if cache is None:
         return whois_func
 
-    async def wrapped(query):
+    async def wrapped(query: str) -> str:
         response = cache.get(query)
         if response is None:
             response = await whois_func(query)
@@ -60,40 +94,41 @@ class LFU:
     2-tuples consisting of a counter giving the number of times this item
     occurs on the eviction queue and the value.
 
+    Args:
+        max_size: Maximum number of entries the cache can contain.
+        max_age: Maximum number of seconds to consider an entry live.
     """
 
     # I may end up reimplementing an LRU cache if it turns out that's more apt,
     # but I haven't went that route as an LRU cache is somewhat more awkward
     # and involved to implement correctly.
 
-    __slots__ = ("cache", "max_age", "max_size", "queue")
+    __slots__ = (
+        "cache",
+        "max_age",
+        "max_size",
+        "queue",
+    )
 
     clock = staticmethod(time.time)
 
-    def __init__(self, max_size=256, max_age=300):
-        """
-        Create a new LFU cache.
-
-        :param max_size int: Maximum number of entries the cache can contain.
-        :param max_age int:  Maximum number of seconds to consider an entry
-                             live.
-        """
+    def __init__(self, max_size: int = 256, max_age: int = 300) -> None:
         super().__init__()
-        self.cache = {}
-        self.queue = collections.deque()
+        self.cache: dict[str, tuple[int, str]] = {}
+        self.queue: t.Deque[tuple[int, str]] = collections.deque()
         self.max_size = int(max_size)
         self.max_age = int(max_age)
 
-    def evict_one(self):
-        """
-        Remove the item at the head of the eviction cache.
-        """
+    def evict_one(self) -> None:
+        """Remove the item at the head of the eviction cache."""
         _, key = self.queue.popleft()
         self.attempt_eviction(key)
 
-    def attempt_eviction(self, key):
-        """
-        Attempt to remove the named item from the cache.
+    def attempt_eviction(self, key: str) -> None:
+        """Attempt to remove the named item from the cache.
+
+        Args:
+            key: The cache key to evict.
         """
         counter, value = self.cache[key]
         counter -= 1
@@ -102,10 +137,8 @@ def attempt_eviction(self, key):
         else:
             self.cache[key] = (counter, value)
 
-    def evict_expired(self):
-        """
-        Evict any items older than the maximum age from the cache.
-        """
+    def evict_expired(self) -> None:
+        """Evict any items older than the maximum age from the cache."""
         cutoff = self.clock() - self.max_age
         while len(self.queue) > 0:
             ts, key = self.queue.popleft()
@@ -114,11 +147,14 @@ def evict_expired(self):
                 break
             self.attempt_eviction(key)
 
-    def get(self, key):
-        """
-        Pull a value from the cache corresponding to the key.
+    def get(self, key: str) -> t.Optional[str]:
+        """Pull a value from the cache corresponding to the key.
 
-        If no value exists, `None` is returned.
+        Args:
+            key: The cache key to look up.
+
+        Returns:
+            The cached value, or `None` if not found.
         """
         self.evict_expired()
         if key not in self.cache:
@@ -128,9 +164,12 @@ def get(self, key):
         self.set(key, value)
         return value
 
-    def set(self, key, value):
-        """
-        Add `value` to the cache, to be referenced by `key`.
+    def set(self, key: str, value: str) -> None:
+        """Add `value` to the cache, to be referenced by `key`.
+
+        Args:
+            key: The cache key to store the value under.
+            value: The value to store.
         """
         if len(self.queue) == self.max_size:
             self.evict_one()
@@ -139,4 +178,4 @@ def set(self, key, value):
         else:
             counter = 0
         self.cache[key] = (counter + 1, value)
-        self.queue.append((self.clock(), key))
+        self.queue.append((int(self.clock()), key))

src/uwhoisd/client.py

@@ -4,6 +4,16 @@
 
 
 async def query_whois(host: str, port: int, query: str) -> str:
+    """Query a WHOIS server.
+
+    Args:
+        host: The WHOIS server hostname.
+        port: The WHOIS server port.
+        query: The WHOIS query.
+
+    Returns:
+        The WHOIS response.
+    """
     reader, writer = await asyncio.open_connection(host, port)
 
     writer.write(f"{query}\r\n".encode())