Add type annotations and improve type checking throughout codebase

.pre-commit-config.yaml

@@ -23,11 +23,16 @@ repos:
     rev: v1.10.0
     hooks:
       - id: mypy
-        additional_dependencies: []
+        # types-PyYAML: mypy reports missing PyYAML stubs even with
+        # ignore_missing_imports, because typeshed tracks them separately.
+        additional_dependencies: [types-PyYAML]
   - repo: local
     hooks:
       - id: pytest
         name: pytest
-        entry: pytest -q
+        # python -m pytest keeps the hook on the interpreter the package is
+        # installed into; a bare ``pytest`` can resolve to an unrelated
+        # standalone install on PATH.
+        entry: python -m pytest -q
         language: system
         pass_filenames: false

pyisolate/__init__.py

@@ -42,10 +42,10 @@
     if isinstance(exc, ImportError) and "cryptography" not in str(exc):
         raise
 
-    def checkpoint(*args, **kwargs):  # type: ignore[no-redef]
+    def checkpoint(sandbox: "Sandbox", key: bytes) -> bytes:  # type: ignore[no-redef]
         raise ModuleNotFoundError("cryptography is required for checkpoint support")
 
-    def restore(*args, **kwargs):  # type: ignore[no-redef]
+    def restore(blob: bytes, key: bytes) -> "Sandbox":  # type: ignore[no-redef]
         raise ModuleNotFoundError("cryptography is required for checkpoint support")
 
 
@@ -91,7 +91,9 @@ def restore(*args, **kwargs):  # type: ignore[no-redef]
     if isinstance(exc, ImportError) and "cryptography" not in str(exc):
         raise
 
-    def migrate(*args, **kwargs):  # type: ignore[no-redef]
+    def migrate(  # type: ignore[no-redef]
+        sandbox: "Sandbox", host: str, key: bytes
+    ) -> "MigrationResponse":
         raise ModuleNotFoundError("cryptography is required for migration support")
 
 

pyisolate/broker/crypto.py

@@ -59,6 +59,9 @@ class CryptoBroker:
     threads; counter and key updates are protected by an internal lock.
     """
 
+    _tx_ctr: int
+    _rx_ctr: int
+
     def __init__(
         self,
         private_key: bytes,

pyisolate/broker/uring.py

@@ -9,6 +9,7 @@
 
 import asyncio
 import os
+from typing import Any
 
 try:
     import uring  # type: ignore
@@ -20,7 +21,7 @@ class IOUring:
     """Minimal wrapper around ``io_uring`` for async file descriptor I/O."""
 
     def __init__(self, entries: int = 8):
-        self._ring = None
+        self._ring: Any = None
         if uring is not None:
             ring = uring.io_uring()
             # handle API differences across wrapper versions

pyisolate/capabilities.py

@@ -40,7 +40,7 @@ class Authority(Capability[str]):
 
     kind: str
 
-    def to_policy_rule(self) -> dict[str, Any]:
+    def to_policy_rule(self) -> dict[str, Any] | str:
         """Return a YAML-serializable representation of this authority."""
         raise NotImplementedError
 

pyisolate/conformance.py

@@ -9,12 +9,12 @@
 
 import argparse
 import json
-import os
 import platform
 import shutil
 import subprocess
 import sys
 import sysconfig
+import threading
 import time
 import uuid
 from dataclasses import asdict, dataclass
@@ -348,8 +348,10 @@ def _probe_cgroup_behavior(self) -> ProbeResult:
         if created:
             cgroup.attach_current(cg_path)
             threads_file = Path(cg_path) / "cgroup.threads"
+            # attach_current writes threading.get_native_id(), so read back the
+            # same identifier; os.gettid() is unavailable on some libc builds.
             attached = threads_file.exists() and str(
-                os.gettid()
+                threading.get_native_id()
             ) in threads_file.read_text(encoding="utf-8")
             cgroup.delete(cg_path)
             deleted = not Path(cg_path).exists()

pyisolate/nogil.py

@@ -50,7 +50,7 @@ def _module_origin(module: ModuleType) -> str | None:
 
 
 def _is_native_origin(origin: str | None) -> bool:
-    return bool(origin) and origin.endswith(_NATIVE_SUFFIXES)
+    return origin is not None and origin.endswith(_NATIVE_SUFFIXES)
 
 
 def imported_native_extensions() -> list[dict[str, Any]]:

pyisolate/observability/alerts.py

@@ -1,15 +1,18 @@
 import logging
+from collections.abc import Callable
 
 logger = logging.getLogger(__name__)
 
+AlertCallback = Callable[[str, Exception], object]
+
 
 class AlertManager:
     """Dispatch callbacks on policy violations."""
 
     def __init__(self) -> None:
-        self._subs: list[callable] = []
+        self._subs: list[AlertCallback] = []
 
-    def register(self, callback) -> None:
+    def register(self, callback: AlertCallback) -> None:
         self._subs.append(callback)
 
     def notify(self, sandbox: str, error: Exception) -> list[Exception]:

pyisolate/operator/__init__.py

@@ -2,10 +2,13 @@
 from __future__ import annotations
 
 import logging
-from typing import Dict
+from typing import TYPE_CHECKING, Dict
 
 from ..supervisor import Supervisor
 
+if TYPE_CHECKING:
+    from ..supervisor import Sandbox
+
 __all__ = ["run_operator", "scale_sandboxes"]
 
 logger = logging.getLogger(__name__)
@@ -19,7 +22,7 @@ def run_operator(namespace: str = "default") -> None:
     api = client.CustomObjectsApi()
     w = watch.Watch()
     sup = Supervisor()
-    sandboxes: Dict[str, object] = {}
+    sandboxes: Dict[str, "Sandbox"] = {}
     for event in w.stream(
         api.list_namespaced_custom_object,
         group="pyisolate.dev",

pyisolate/policy/__init__.py

@@ -5,6 +5,7 @@
 import socket
 import tempfile
 import urllib.request
+from collections.abc import Iterable, Mapping
 from dataclasses import dataclass, field
 from importlib import resources
 from pathlib import Path
@@ -42,8 +43,8 @@ def _mini_load(text: str) -> dict:
                     from typing import cast
 
                     assert isinstance(result[current], list)
-                    lst = cast(list[str], result[current])
-                    lst.append(_unquote(item))
+                    str_lst = cast(list[str], result[current])
+                    str_lst.append(_unquote(item))
                 continue
 
             if ":" not in line:
@@ -66,7 +67,7 @@ def safe_load(stream):
                 return _mini_load(stream.read())
             return _mini_load(stream)
 
-    yaml = _MiniYaml()
+    yaml = _MiniYaml()  # type: ignore[assignment]
 
 
 from ..capabilities import ConnectTCP, CpuBudget, Import, ReadPath, WritePath
@@ -188,11 +189,13 @@ def to_yaml(self, name: str = "default") -> str:
         lines = ["version: 1.0", "sandboxes:", f"  {name}:"]
         sandbox = data["sandboxes"][name]  # type: ignore[index]
 
-        def append_rule_list(section: str, rules: object) -> None:
+        def append_rule_list(
+            section: str, rules: Iterable[Mapping[str, object]] | None
+        ) -> None:
             if not rules:
                 return
             lines.append(f"    {section}:")
-            for rule in rules:  # type: ignore[union-attr]
+            for rule in rules:
                 key, value = next(iter(rule.items()))
                 lines.append(f"      - {key}: {value!r}")
 

pyisolate/policy/compiler.py

@@ -299,10 +299,10 @@ def compile_policy(path: str | Path) -> CompiledPolicy:
                 capabilities.append(ReadPath(rule.path))
             elif rule.action == "write":
                 capabilities.append(WritePath(rule.path))
-        for rule in tcp_compiled:
-            if rule.action == "connect":
+        for tcp_rule in tcp_compiled:
+            if tcp_rule.action == "connect":
                 try:
-                    capabilities.append(ConnectTCP.from_address(rule.addr))
+                    capabilities.append(ConnectTCP.from_address(tcp_rule.addr))
                 except ValueError:
                     pass
         capabilities.extend(Import(module) for module in imports)

pyisolate/policy/model.py

@@ -280,11 +280,11 @@ def from_sandbox_policy(policy: Any) -> RuntimePolicy:
         if isinstance(item, str):
             allow_tcp.append(NetworkRule("connect", item))
             continue
-        rule = _coerce_compiled_tcp(item)
-        if rule.action == "connect":
-            allow_tcp.append(rule)
+        tcp_rule = _coerce_compiled_tcp(item)
+        if tcp_rule.action == "connect":
+            allow_tcp.append(tcp_rule)
         else:
-            deny_tcp.append(rule)
+            deny_tcp.append(tcp_rule)
 
     imports = tuple(str(module) for module in (getattr(policy, "imports", None) or ()))
     return RuntimePolicy(
@@ -382,13 +382,13 @@ def to_bpf_map_entries(policy_set: RuntimePolicySet) -> list[tuple[str, str, str
             entries.append(("policy_fs_allow", f"{sandbox_name}:{index}", rule.path))
         for index, rule in enumerate(policy.deny_fs):
             entries.append(("policy_fs_deny", f"{sandbox_name}:{index}", rule.path))
-        for index, rule in enumerate(policy.allow_tcp):
+        for index, net_rule in enumerate(policy.allow_tcp):
             entries.append(
-                ("policy_net_allow", f"{sandbox_name}:{index}", rule.destination)
+                ("policy_net_allow", f"{sandbox_name}:{index}", net_rule.destination)
             )
-        for index, rule in enumerate(policy.deny_tcp):
+        for index, net_rule in enumerate(policy.deny_tcp):
             entries.append(
-                ("policy_net_deny", f"{sandbox_name}:{index}", rule.destination)
+                ("policy_net_deny", f"{sandbox_name}:{index}", net_rule.destination)
             )
         for index, module in enumerate(policy.imports):
             entries.append(("policy_import_allow", f"{sandbox_name}:{index}", module))

pyisolate/runtime/thread.py

@@ -48,7 +48,7 @@
 from ..numa import bind_current_thread
 from ..observability.trace import Tracer
 from ..policy.model import RuntimePolicy, from_sandbox_policy
-from ..telemetry import DenialEvent
+from ..telemetry import Decision, DenialEvent
 from .protocol import (
     AttachCgroupRequest,
     BrokerRequest,
@@ -123,8 +123,8 @@ def _deny(
     policy_rule: str,
     message: str,
     *,
-    kernel_decision: str = "not_evaluated",
-    broker_decision: str = "deny",
+    kernel_decision: Decision = "not_evaluated",
+    broker_decision: Decision = "deny",
 ) -> errors.PolicyError:
     sandbox = _active_sandbox()
     cell = sandbox.name if sandbox is not None else "<unknown>"
@@ -670,12 +670,12 @@ def _check_network_destination(address: Iterable[str]) -> None:
             raise errors.NetworkExceeded()
 
 
-def _guarded_connect(self_socket: socket.socket, address: Iterable[str]):
+def _guarded_connect(self_socket: socket.socket, address: Any):
     _check_network_destination(address)
     return _ORIG_SOCKET_CONNECT(self_socket, address)
 
 
-def _guarded_connect_ex(self_socket: socket.socket, address: Iterable[str]):
+def _guarded_connect_ex(self_socket: socket.socket, address: Any):
     _check_network_destination(address)
     return _ORIG_SOCKET_CONNECT_EX(self_socket, address)
 
@@ -771,7 +771,7 @@ def _run_with_accounting(*r_args, **r_kwargs):
     return SandboxedThread
 
 
-def _module_proxy(name: str, module) -> types.ModuleType:
+def _module_proxy(name: str, module) -> Any:
     proxy = types.ModuleType(name, getattr(module, "__doc__", None))
     proxy.__dict__.update({k: getattr(module, k) for k in dir(module)})
     proxy.__dict__["__name__"] = name
@@ -840,7 +840,7 @@ def _wrap_module(name: str, module):
         return mod
     if base == "time":
 
-        def _require_clock() -> ClockCapability:
+        def _require_clock() -> ClockCapability | None:
             cap = getattr(_thread_local, "clock_capability", None)
             return cap
 
@@ -885,7 +885,8 @@ def __init__(self, family=-1, type=-1, proto=-1, fileno=None):
 
             connect = _guarded_connect
             connect_ex = _guarded_connect_ex
-            sendto = _guarded_sendto
+            # The guard collapses socket.sendto's overloads into one signature.
+            sendto = _guarded_sendto  # type: ignore[assignment]
 
         def _create_connection(
             address, timeout=socket._GLOBAL_DEFAULT_TIMEOUT, source_address=None
@@ -936,7 +937,7 @@ def _create_connection(
     if base == "pathlib":
         mod = _module_proxy("pathlib", module)
 
-        class SandboxedPath(type(module.Path())):
+        class SandboxedPath(type(module.Path())):  # type: ignore[misc]
             def open(
                 self,
                 mode="r",
@@ -1048,6 +1049,12 @@ class Stats:
 class SandboxThread(threading.Thread):
     """Thread that runs guest code and communicates via a queue."""
 
+    # Assigned by the supervisor after construction/reset rather than in
+    # ``__init__``; declared here so the attributes are part of the class
+    # contract.
+    _backend: str
+    _temp_dir: Path
+
     @staticmethod
     def _merge_allowed_imports(policy, allowed_imports: Optional[Iterable[str]]):
         imports: set[str] = set()
@@ -1066,7 +1073,7 @@ def _merge_allowed_imports(policy, allowed_imports: Optional[Iterable[str]]):
             imports.update(runtime_policy.imports)
 
         allowed_imports_provided = allowed_imports is not None
-        if allowed_imports_provided:
+        if allowed_imports is not None:
             imports.update(allowed_imports)
 
         if not imports and not allowed_imports_provided:
@@ -1116,7 +1123,7 @@ def _init_config_wiring(
         self.child_work_max = child_work_max
         self.allowed_imports = self._merge_allowed_imports(policy, allowed_imports)
         self.numa_node = numa_node
-        self._bound_numa_node = None
+        self._bound_numa_node: Optional[int] = None
         self._cgroup_path = cgroup_path
         self.quota_enforcement = enforcement_status
         self._capabilities = deserialize_capabilities(capabilities)
@@ -1145,15 +1152,15 @@ def _reset_runtime_state(self) -> None:
         self._cpu_time = 0.0
         self._mem_peak = 0
         self._mem_base = 0
-        self._start_time = None
+        self._start_time: Optional[float] = None
         self._ops = 0
         self._errors = 0
         self._latency = {"0.5": 0, "1": 0, "5": 0, "10": 0, "inf": 0}
         self._latency_sum = 0.0
         self._trace_enabled = False
         self._syscall_log: list[str] = []
-        self._quarantine_reason = None
-        self.termination_reason = None
+        self._quarantine_reason: Optional[str] = None
+        self.termination_reason: Optional[str] = None
         self._open_files = 0
         self._network_ops = 0
         self._output_bytes = 0
@@ -1172,7 +1179,7 @@ def __init__(
         output_bytes_max: Optional[int] = None,
         child_work_max: Optional[int] = None,
         allowed_imports: Optional[list[str]] = None,
-        on_violation: Optional[Callable[[str, Exception], None]] = None,
+        on_violation: Optional[Callable[[str, Exception], object]] = None,
         tracer: Optional["Tracer"] = None,
         numa_node: Optional[int] = None,
         cgroup_path=None,