chore: add type hints (#465)

Co-authored-by: cbornet <cbornet@hotmail.com>

Author: miketheman

.github/workflows/python.yml

@@ -41,6 +41,8 @@ jobs:
       run: uv python install '${{ matrix.python-version }}' # zizmor: ignore[template-injection]
     - name: Prepare project for development
       run: uv sync
+    - name: Check types
+      run: uv run mypy src
     - name: Test with pytest
       run: |
         uv run coverage run -m pytest

pyproject.toml

@@ -20,6 +20,7 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: CPython",
     "Operating System :: OS Independent",
     "License :: OSI Approved :: MIT License",
+    "Typing :: Typed",
 ]
 requires-python = ">= 3.10"
 dependencies = ["pytest>=7.0.0"]
@@ -32,6 +33,7 @@ dev = [
     "requests >= 2.32.4",
     "starlette >= 0.47.1",
     "httpx >= 0.28.1",
+    "mypy >= 1.20",
 ]
 
 [project.urls]
@@ -55,6 +57,9 @@ known_first_party = ['pytest_socket', 'conftest']
 # https://black.readthedocs.io/en/stable/guides/using_black_with_other_tools.html#profilegcm
 profile = "black"
 
+[tool.mypy]
+strict = true
+
 [tool.vulture]
 ignore_decorators = ["@pytest.fixture"]
 ignore_names = ["pytest_*"]

src/pytest_socket/__init__.py

@@ -1,8 +1,12 @@
+from __future__ import annotations
+
 import ipaddress
 import itertools
 import socket
 from collections import defaultdict
+from collections.abc import Iterator
 from dataclasses import dataclass, field
+from typing import Any
 
 import pytest
 
@@ -11,21 +15,26 @@
 
 
 class SocketBlockedError(RuntimeError):
-    def __init__(self, *_args, **_kwargs):
+    def __init__(self, *_args: Any, **_kwargs: Any) -> None:
         super().__init__("A test tried to use socket.socket.")
 
 
 class SocketConnectBlockedError(RuntimeError):
-    def __init__(self, allowed, host, *_args, **_kwargs):
-        if allowed:
-            allowed = ",".join(allowed)
+    def __init__(
+        self,
+        allowed: list[str],
+        host: str | None,
+        *_args: Any,
+        **_kwargs: Any,
+    ) -> None:
+        allowed_str = ",".join(allowed)
         super().__init__(
             "A test tried to use socket.socket.connect() "
-            f'with host "{host}" (allowed: "{allowed}").'
+            f'with host "{host}" (allowed: "{allowed_str}").'
         )
 
 
-def pytest_addoption(parser):
+def pytest_addoption(parser: pytest.Parser) -> None:
     group = parser.getgroup("socket")
     group.addoption(
         "--disable-socket",
@@ -54,15 +63,15 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture
-def socket_disabled(pytestconfig):
+def socket_disabled(pytestconfig: pytest.Config) -> Iterator[None]:
     """disable socket.socket for duration of this test function"""
     socket_config = pytestconfig.stash[_STASH_KEY]
     disable_socket(allow_unix_socket=socket_config.allow_unix_socket)
     yield
 
 
 @pytest.fixture
-def socket_enabled(pytestconfig):
+def socket_enabled(pytestconfig: pytest.Config) -> Iterator[None]:
     """enable socket.socket for duration of this test function"""
     enable_socket()
     yield
@@ -80,31 +89,37 @@ class _PytestSocketConfig:
 _STASH_KEY = pytest.StashKey[_PytestSocketConfig]()
 
 
-def _is_unix_socket(family) -> bool:
+def _is_unix_socket(family: int) -> bool:
     return hasattr(socket, "AF_UNIX") and family == socket.AF_UNIX
 
 
-def disable_socket(allow_unix_socket=False):
+def disable_socket(allow_unix_socket: bool = False) -> None:
     """disable socket.socket to disable the Internet. useful in testing."""
 
     class GuardedSocket(socket.socket):
         """socket guard to disable socket creation (from pytest-socket)"""
 
-        def __new__(cls, family=-1, type=-1, proto=-1, fileno=None):
+        def __new__(
+            cls,
+            family: socket.AddressFamily | int = -1,
+            type: socket.SocketKind | int = -1,
+            proto: int = -1,
+            fileno: int | None = None,
+        ) -> GuardedSocket:
             if _is_unix_socket(family) and allow_unix_socket:
-                return super().__new__(cls, family, type, proto, fileno)
+                return super().__new__(cls, family, type, proto, fileno)  # type: ignore[call-arg] # noqa E501
 
             raise SocketBlockedError()
 
-    socket.socket = GuardedSocket
+    socket.socket = GuardedSocket  # type: ignore[misc]
 
 
-def enable_socket():
+def enable_socket() -> None:
     """re-enable socket.socket to enable the Internet. useful in testing."""
-    socket.socket = _true_socket
+    socket.socket = _true_socket  # type: ignore[misc]
 
 
-def pytest_configure(config):
+def pytest_configure(config: pytest.Config) -> None:
     config.addinivalue_line(
         "markers", "disable_socket(): Disable socket connections for a specific test"
     )
@@ -125,7 +140,7 @@ def pytest_configure(config):
     )
 
 
-def pytest_runtest_setup(item) -> None:
+def pytest_runtest_setup(item: pytest.Item) -> None:
     """During each test item's setup phase,
     choose the behavior based on the configurations supplied.
 
@@ -166,7 +181,7 @@ def pytest_runtest_setup(item) -> None:
         disable_socket(socket_config.allow_unix_socket)
 
 
-def _resolve_allow_hosts(item):
+def _resolve_allow_hosts(item: pytest.Item) -> str | list[str] | None:
     """Resolve `allow_hosts` behaviors."""
     socket_config = item.config.stash[_STASH_KEY]
 
@@ -186,21 +201,23 @@ def _resolve_allow_hosts(item):
     return hosts
 
 
-def pytest_runtest_teardown():
+def pytest_runtest_teardown() -> None:
     _remove_restrictions()
 
 
-def host_from_address(address):
+def host_from_address(address: tuple[Any, ...]) -> str | None:
     host = address[0]
     if isinstance(host, str):
         return host
+    return None
 
 
-def host_from_connect_args(args):
+def host_from_connect_args(args: tuple[Any, ...]) -> str | None:
     address = args[0]
 
     if isinstance(address, tuple):
         return host_from_address(address)
+    return None
 
 
 def is_ipaddress(address: str) -> bool:
@@ -217,15 +234,16 @@ def is_ipaddress(address: str) -> bool:
 def resolve_hostnames(hostname: str) -> set[str]:
     try:
         return {
-            addr_struct[0] for *_, addr_struct in socket.getaddrinfo(hostname, None)
+            addr_struct[0]  # type: ignore[misc]
+            for *_, addr_struct in socket.getaddrinfo(hostname, None)
         }
     except socket.gaierror:
         return set()
 
 
 def normalize_allowed_hosts(
     allowed_hosts: list[str],
-    resolution_cache: dict[str, list[str]] | None = None,
+    resolution_cache: dict[str, set[str]] | None = None,
 ) -> dict[str, set[str]]:
     """Map all items in `allowed_hosts` to IP addresses."""
     if resolution_cache is None:
@@ -246,7 +264,7 @@ def normalize_allowed_hosts(
 def socket_allow_hosts(
     allowed: str | list[str] | None = None,
     allow_unix_socket: bool = False,
-    resolution_cache: dict[str, list[str]] | None = None,
+    resolution_cache: dict[str, set[str]] | None = None,
 ) -> None:
     """disable socket.socket.connect() to disable the Internet. useful in testing."""
     if isinstance(allowed, str):
@@ -270,7 +288,7 @@ def socket_allow_hosts(
         ]
     )
 
-    def guarded_connect(inst, *args):
+    def guarded_connect(inst: socket.socket, *args: Any) -> None:
         host = host_from_connect_args(args)
         if host in allowed_ip_hosts_and_hostnames or (
             _is_unix_socket(inst.family) and allow_unix_socket
@@ -279,10 +297,10 @@ def guarded_connect(inst, *args):
 
         raise SocketConnectBlockedError(allowed_list, host)
 
-    socket.socket.connect = guarded_connect
+    socket.socket.connect = guarded_connect  # type: ignore[assignment,method-assign]
 
 
-def _remove_restrictions():
+def _remove_restrictions() -> None:
     """restore socket.socket.* to allow access to the Internet. useful in testing."""
-    socket.socket = _true_socket
-    socket.socket.connect = _true_connect
+    socket.socket = _true_socket  # type: ignore[misc]
+    socket.socket.connect = _true_connect  # type: ignore[method-assign]