Add strict type annotations

Author: asvetlov

CHANGES.rst

@@ -7,6 +7,8 @@
 
 * Provide ``X-Forwarded`` middleware that filters out trusted values (#153)
 
+* Add type annotations
+
 0.1.2 (2018-03-01)
 ==================
 

Makefile

@@ -11,7 +11,7 @@ endif
 
 .PHONY: lint
 lint: fmt
-	mypy aiohttp_remotes
+	mypy --strict --show-error-codes aiohttp_remotes tests
 
 test:
 	pytest tests

aiohttp_remotes/__init__.py

@@ -9,6 +9,10 @@
 __version__ = "1.0.0a0"
 
 
+from typing_extensions import Protocol
+
+from aiohttp import web
+
 from .allowed_hosts import AllowedHosts
 from .basic_auth import BasicAuth
 from .cloudflare import Cloudflare
@@ -17,7 +21,12 @@
 from .x_forwarded import XForwardedFiltered, XForwardedRelaxed, XForwardedStrict
 
 
-async def setup(app, *tools):
+class _Tool(Protocol):
+    async def setup(self, app: web.Application) -> None:
+        ...
+
+
+async def setup(app: web.Application, *tools: _Tool) -> None:
     for tool in tools:
         await tool.setup(app)
 

aiohttp_remotes/abc.py

@@ -1,12 +1,14 @@
 import abc
 
+from typing_extensions import NoReturn
+
 from aiohttp import web
 
 
 class ABC(abc.ABC):
     @abc.abstractmethod
-    async def setup(self, app):
+    async def setup(self, app: web.Application) -> None:
         pass  # pragma: no cover
 
-    async def raise_error(self, request):
+    async def raise_error(self, request: web.Request) -> NoReturn:
         raise web.HTTPBadRequest()

aiohttp_remotes/allowed_hosts.py

@@ -1,32 +1,43 @@
+from typing import Awaitable, Callable, Iterable, Set, Union
+
 from aiohttp import web
 
 from .abc import ABC
 
 
 class ANY:
-    def __contains__(self, item):
+    def __contains__(self, item: object) -> bool:
         return True
 
 
 class AllowedHosts(ABC):
-    def __init__(self, allowed_hosts=("*",), *, white_paths=()):
-        allowed_hosts = set(allowed_hosts)
-
-        if "*" in allowed_hosts:
-            allowed_hosts = ANY()
-
-        self._allowed_hosts = allowed_hosts
+    def __init__(
+        self,
+        allowed_hosts: Iterable[str] = ("*",),
+        *,
+        white_paths: Iterable[str] = (),
+    ) -> None:
+        real_allowed_hosts: Union[Set[str], ANY] = set(allowed_hosts)
+
+        if "*" in real_allowed_hosts:
+            real_allowed_hosts = ANY()
+
+        self._allowed_hosts = real_allowed_hosts
         self._white_paths = set(white_paths)
 
-    async def setup(self, app):
+    async def setup(self, app: web.Application) -> None:
         app.middlewares.append(self.middleware)
 
     @web.middleware
-    async def middleware(self, request, handler):
+    async def middleware(
+        self,
+        request: web.Request,
+        handler: Callable[[web.Request], Awaitable[web.StreamResponse]],
+    ) -> web.StreamResponse:
         if (
             request.path not in self._white_paths
             and request.host not in self._allowed_hosts
         ):
-            await self.raise_error(request)
+            return await self.raise_error(request)
 
         return await handler(request)

aiohttp_remotes/basic_auth.py

@@ -1,28 +1,43 @@
 import base64
 import binascii
+from typing import Awaitable, Callable, Iterable
+
+from typing_extensions import NoReturn
 
 from aiohttp import hdrs, web
 
 from .abc import ABC
 
 
 class BasicAuth(ABC):
-    def __init__(self, username, password, realm, *, white_paths=()):
+    def __init__(
+        self,
+        username: str,
+        password: str,
+        realm: str,
+        *,
+        white_paths: Iterable[str] = (),
+    ) -> None:
         self._username = username
         self._password = password
         self._realm = realm
         self._white_paths = set(white_paths)
 
-    async def setup(self, app):
+    async def setup(self, app: web.Application) -> None:
         app.middlewares.append(self.middleware)
 
-    async def raise_error(self, request):
+    async def raise_error(self, request: web.Request) -> NoReturn:
         raise web.HTTPUnauthorized(
             headers={hdrs.WWW_AUTHENTICATE: f"Basic realm={self._realm}"},
         )
 
     @web.middleware
-    async def middleware(self, request, handler):
+    async def middleware(
+        self,
+        request: web.Request,
+        handler: Callable[[web.Request], Awaitable[web.StreamResponse]],
+    ) -> web.StreamResponse:
+
         if request.path not in self._white_paths:
             auth_header = request.headers.get(hdrs.AUTHORIZATION)
 
@@ -34,16 +49,16 @@ async def middleware(self, request, handler):
 
                 auth_decoded = base64.decodebytes(secret).decode("utf-8")
             except (UnicodeDecodeError, UnicodeEncodeError, binascii.Error):
-                await self.raise_error(request)
+                return await self.raise_error(request)
 
             credentials = auth_decoded.split(":")
 
             if len(credentials) != 2:
-                await self.raise_error(request)
+                return await self.raise_error(request)
 
             username, password = credentials
 
             if username != self._username or password != self._password:
-                await self.raise_error(request)
+                return await self.raise_error(request)
 
         return await handler(request)

aiohttp_remotes/cloudflare.py

@@ -1,29 +1,31 @@
 from ipaddress import ip_address, ip_network
+from typing import Awaitable, Callable, Optional, Set
 
 import aiohttp
 from aiohttp import web
 
 from .abc import ABC
+from .exceptions import IPNetwork
 from .log import logger
 
 
 class Cloudflare(ABC):
-    def __init__(self, client=None):
-        self._ip_networks = set()
+    def __init__(self, client: Optional[aiohttp.ClientSession] = None) -> None:
+        self._ip_networks: Set[IPNetwork] = set()
         self._client = client
 
-    def _parse_mask(self, text):
+    def _parse_mask(self, text: str) -> Set[IPNetwork]:
         ret = set()
         for mask in text.splitlines():
             try:
-                mask = ip_network(mask)
+                real_mask = ip_network(mask)
             except (ValueError, TypeError):
                 continue
 
-            ret.add(mask)
+            ret.add(real_mask)
         return ret
 
-    async def setup(self, app):
+    async def setup(self, app: web.Application) -> None:
         if self._client is not None:  # pragma: no branch
             client = self._client
         else:
@@ -42,7 +44,11 @@ async def setup(self, app):
         app.middlewares.append(self.middleware)
 
     @web.middleware
-    async def middleware(self, request, handler):
+    async def middleware(
+        self,
+        request: web.Request,
+        handler: Callable[[web.Request], Awaitable[web.StreamResponse]],
+    ) -> web.StreamResponse:
         remote_ip = ip_address(request.remote)
 
         for network in self._ip_networks:
@@ -54,4 +60,4 @@ async def middleware(self, request, handler):
         context = {"remote_ip": remote_ip}
         logger.error(msg, context)
 
-        await self.raise_error(request)
+        return await self.raise_error(request)

aiohttp_remotes/exceptions.py

@@ -1,87 +1,98 @@
+import builtins
+from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
+from typing import Any, Dict, Sequence, Union, cast
+
+from aiohttp import web
+
 from .log import logger
 
+IPAddress = Union[IPv4Address, IPv6Address]
+IPNetwork = Union[IPv4Network, IPv6Network]
+IPRule = Union[IPAddress, IPNetwork]
+Trusted = Sequence[Union["builtins.ellipsis", Sequence[IPRule]]]
+
 
 class RemoteError(Exception):
-    def log(self, request):
+    def log(self, request: web.Request) -> None:
         raise NotImplementedError  # pragma: no cover
 
 
 class TooManyHeaders(RemoteError):
     @property
-    def header(self):
-        return self.args[0]
+    def header(self) -> str:
+        return cast(str, self.args[0])
 
-    def log(self, request):
+    def log(self, request: web.Request) -> None:
         msg = "Too many headers for %(header)s"
-        context = {"header": self.header}
+        context: Dict[str, Any] = {"header": self.header}
         extra = context.copy()
         extra["request"] = request
         logger.error(msg, context, extra=extra)
 
 
 class IncorrectIPCount(RemoteError):
     @property
-    def expected(self):
-        return self.args[0]
+    def expected(self) -> int:
+        return cast(int, self.args[0])
 
     @property
-    def actual(self):
-        return self.args[1]
+    def actual(self) -> Sequence[IPAddress]:
+        return cast(Sequence[IPAddress], self.args[1])
 
-    def log(self, request):
+    def log(self, request: web.Request) -> None:
         msg = "Too many X-Forwarded-For values: %(actual)s, " "expected %(expected)s"
-        context = {"actual": self.actual, "expected": self.expected}
+        context: Dict[str, Any] = {"actual": self.actual, "expected": self.expected}
         extra = context.copy()
         extra["request"] = request
         logger.error(msg, context, extra=extra)
 
 
 class IncorrectForwardedCount(RemoteError):
     @property
-    def expected(self):
-        return self.args[0]
+    def expected(self) -> int:
+        return cast(int, self.args[0])
 
     @property
-    def actual(self):
-        return self.args[1]
+    def actual(self) -> int:
+        return cast(int, self.args[1])
 
-    def log(self, request):
+    def log(self, request: web.Request) -> None:
         msg = "Too many Forwarded values: %(actual)s, " "expected %(expected)s"
-        context = {"actual": self.actual, "expected": self.expected}
+        context: Dict[str, Any] = {"actual": self.actual, "expected": self.expected}
         extra = context.copy()
         extra["request"] = request
         logger.error(msg, context, extra=extra)
 
 
 class IncorrectProtoCount(RemoteError):
     @property
-    def expected(self):
-        return self.args[0]
+    def expected(self) -> int:
+        return cast(int, self.args[0])
 
     @property
-    def actual(self):
-        return self.args[1]
+    def actual(self) -> Sequence[str]:
+        return cast(Sequence[str], self.args[1])
 
-    def log(self, request):
+    def log(self, request: web.Request) -> None:
         msg = "Too many X-Forwarded-Proto values: %(actual)s, " "expected %(expected)s"
-        context = {"actual": self.actual, "expected": self.expected}
+        context: Dict[str, Any] = {"actual": self.actual, "expected": self.expected}
         extra = context.copy()
         extra["request"] = request
         logger.error(msg, context, extra=extra)
 
 
 class UntrustedIP(RemoteError):
     @property
-    def ip(self):
-        return self.args[0]
+    def ip(self) -> IPAddress:
+        return cast(IPAddress, self.args[0])
 
     @property
-    def trusted(self):
-        return self.args[1]
+    def trusted(self) -> Sequence[IPAddress]:
+        return cast(Sequence[IPAddress], self.args[1])
 
-    def log(self, request):
+    def log(self, request: web.Request) -> None:
         msg = "Untrusted IP: %(ip)s, trusted: %(trusted)s"
-        context = {"ip": self.ip, "trusted": self.trusted}
+        context: Dict[str, Any] = {"ip": self.ip, "trusted": self.trusted}
         extra = context.copy()
         extra["request"] = request
         logger.error(msg, context, extra=extra)