From 7b677e8150e9be071e2256cf3122727a56cad1e0 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:04:08 +0400
Subject: [PATCH 01/10] Annotate the runtime cache and simulation gateway
 clients

Type the cache backend protocol with what callers rely on (string
envelopes in, opaque write acknowledgements out) and add a CachePipeline
protocol for the queued writes. Decoded envelope payloads are object,
because every consumer validates them. A SPMCapabilitySource protocol
covers both the HTTP entrypoint client and the release registry adapter
that validate_worker_spm accepts. ReformImpactCache.update keeps Any for
the keyword changes it forwards unchanged to dataclasses.replace.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 policyengine_api/libs/gateway_auth.py         | 12 +++++--
 .../libs/simulation_entrypoint.py             | 10 ++++--
 policyengine_api/runtime_cache/core.py        | 34 +++++++++++++------
 policyengine_api/runtime_cache/fake.py        | 18 +++++-----
 .../runtime_cache/reform_impacts.py           |  6 ++--
 policyengine_api/worker_spm.py                | 28 ++++++++++++---
 policyengine_api/worker_spm_release.py        | 10 +++---
 7 files changed, 82 insertions(+), 36 deletions(-)

diff --git a/policyengine_api/libs/gateway_auth.py b/policyengine_api/libs/gateway_auth.py
index 808202e5b..8c4b28252 100644
--- a/policyengine_api/libs/gateway_auth.py
+++ b/policyengine_api/libs/gateway_auth.py
@@ -13,9 +13,13 @@
 import threading
 import time
 from functools import cache
+from typing import TYPE_CHECKING
 
 import httpx
 
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
 GATEWAY_AUTH_ISSUER_ENV = "GATEWAY_AUTH_ISSUER"
 GATEWAY_AUTH_AUDIENCE_ENV = "GATEWAY_AUTH_AUDIENCE"
 GATEWAY_AUTH_CLIENT_ID_ENV = "GATEWAY_AUTH_CLIENT_ID"
@@ -115,7 +119,7 @@ def __init__(
         client_secret_resource: str | None = None,
         *,
         http_timeout: float = 10.0,
-    ):
+    ) -> None:
         self._issuer = (
             issuer
             if issuer is not None
@@ -251,10 +255,12 @@ def invalidate(self) -> None:
 class GatewayBearerAuth(httpx.Auth):
     """Attach a bearer token and retry once on a 401."""
 
-    def __init__(self, token_provider: GatewayAuthTokenProvider):
+    def __init__(self, token_provider: GatewayAuthTokenProvider) -> None:
         self._token_provider = token_provider
 
-    def auth_flow(self, request):
+    def auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response, None]:
         request.headers["Authorization"] = f"Bearer {self._token_provider.get_token()}"
         response = yield request
 
diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py
index 997cb0352..9e33a8cfe 100644
--- a/policyengine_api/libs/simulation_entrypoint.py
+++ b/policyengine_api/libs/simulation_entrypoint.py
@@ -118,7 +118,7 @@ class SimulationEntrypointClient:
     polling for their status/results via HTTP endpoints.
     """
 
-    def __init__(self, entrypoint: str | None = None):
+    def __init__(self, entrypoint: str | None = None) -> None:
         self.entrypoint = entrypoint or get_sim_entrypoint()
         self.base_url = resolve_simulation_entrypoint_url(self.entrypoint)
         self._token_provider = GatewayAuthTokenProvider()
@@ -308,7 +308,13 @@ def run_budget_window_batch(self, payload: dict) -> ModalBudgetWindowBatchExecut
             )
             raise
 
-    def get_spm_capability(self, country, version=None, *, policyengine_version=None):
+    def get_spm_capability(
+        self,
+        country: str,
+        version: str | None = None,
+        *,
+        policyengine_version: str | None = None,
+    ) -> object:
         from policyengine_api.spm import SPMValidationError
 
         try:
diff --git a/policyengine_api/runtime_cache/core.py b/policyengine_api/runtime_cache/core.py
index 8b75dabfd..7d50bdfda 100644
--- a/policyengine_api/runtime_cache/core.py
+++ b/policyengine_api/runtime_cache/core.py
@@ -7,7 +7,7 @@
 import random
 import time
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol, Self
 
 from policyengine_api.gcp_logging import logger
 
@@ -18,12 +18,26 @@
 COMPLETED_RESULT_TTL_JITTER_FRACTION = 0.10
 
 
+class CachePipeline(Protocol):
+    """The queued commands that cache writers issue inside one pipeline."""
+
+    def __enter__(self) -> Self: ...
+    def __exit__(self, *args: object) -> object: ...
+    def set(self, key: str, value: str, **kwargs: object) -> object: ...
+    def delete(self, *keys: str) -> object: ...
+    def expire(self, key: str, seconds: int) -> object: ...
+    def zadd(self, key: str, mapping: dict[str, float]) -> object: ...
+    def zrem(self, key: str, *members: str) -> object: ...
+    def zremrangebyrank(self, key: str, start: int, end: int) -> object: ...
+    def execute(self) -> list[object]: ...
+
+
 class CacheBackend(Protocol):
-    def get(self, key: str) -> Any: ...
-    def set(self, key: str, value: str, **kwargs: Any) -> Any: ...
-    def delete(self, *keys: str) -> Any: ...
-    def pipeline(self, transaction: bool = True): ...
-    def eval(self, script: str, numkeys: int, *keys_and_args: str) -> Any: ...
+    def get(self, key: str) -> str | bytes | None: ...
+    def set(self, key: str, value: str, **kwargs: object) -> object: ...
+    def delete(self, *keys: str) -> object: ...
+    def pipeline(self, transaction: bool = True) -> CachePipeline: ...
+    def eval(self, script: str, numkeys: int, *keys_and_args: str) -> object: ...
 
 
 class CacheCoordinationError(RuntimeError):
@@ -66,7 +80,7 @@ def family_key(self, family: str, schema_version: int, name: str) -> str:
 def encode_envelope(
     family: str,
     schema_version: int,
-    payload: Any,
+    payload: object,
 ) -> str:
     return json.dumps(
         {
@@ -85,7 +99,7 @@ def decode_envelope(
     *,
     family: str,
     schema_version: int,
-) -> Any | None:
+) -> object:
     if value is None:
         return None
     try:
@@ -171,7 +185,7 @@ def __init__(
     def key(self, inputs: dict[str, Any]) -> str:
         return self.namespace.key(self.family, self.schema_version, inputs)
 
-    def get(self, inputs: dict[str, Any]) -> Any | None:
+    def get(self, inputs: dict[str, Any]) -> object:
         started_at = time.perf_counter()
         try:
             value = self.client.get(self.key(inputs))
@@ -195,7 +209,7 @@ def get(self, inputs: dict[str, Any]) -> Any | None:
         )
         return payload
 
-    def set(self, inputs: dict[str, Any], payload: Any) -> bool:
+    def set(self, inputs: dict[str, Any], payload: object) -> bool:
         started_at = time.perf_counter()
         try:
             result = self.client.set(
diff --git a/policyengine_api/runtime_cache/fake.py b/policyengine_api/runtime_cache/fake.py
index 1a9008659..74fe06318 100644
--- a/policyengine_api/runtime_cache/fake.py
+++ b/policyengine_api/runtime_cache/fake.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from threading import RLock
-from typing import TYPE_CHECKING, Any, Self
+from typing import TYPE_CHECKING, Self
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -15,7 +15,7 @@ class InMemoryCacheBackend:
     def __init__(self, *, now: Callable[[], float] | None = None) -> None:
         self._external_now = now
         self._time = 0.0
-        self._values: dict[str, Any] = {}
+        self._values: dict[str, str] = {}
         self._expires: dict[str, float] = {}
         self._sorted_sets: dict[str, dict[str, float]] = {}
         self._lock = RLock()
@@ -35,7 +35,7 @@ def _purge(self, key: str) -> None:
             self._sorted_sets.pop(key, None)
             self._expires.pop(key, None)
 
-    def get(self, key: str) -> Any:
+    def get(self, key: str) -> str | None:
         with self._lock:
             self._purge(key)
             return self._values.get(key)
@@ -43,7 +43,7 @@ def get(self, key: str) -> Any:
     def set(
         self,
         key: str,
-        value: Any,
+        value: str,
         *,
         ex: int | None = None,
         nx: bool = False,
@@ -139,16 +139,16 @@ class InMemoryPipeline:
     def __init__(self, backend: InMemoryCacheBackend, *, transaction: bool) -> None:
         self.backend = backend
         self.transaction = transaction
-        self.operations: list[tuple[str, tuple[Any, ...], dict[str, Any]]] = []
+        self.operations: list[tuple[str, tuple[object, ...], dict[str, object]]] = []
 
-    def __getattr__(self, name: str):
-        def queue(*args: Any, **kwargs: Any) -> InMemoryPipeline:
+    def __getattr__(self, name: str) -> Callable[..., InMemoryPipeline]:
+        def queue(*args: object, **kwargs: object) -> InMemoryPipeline:
             self.operations.append((name, args, kwargs))
             return self
 
         return queue
 
-    def execute(self) -> list[Any]:
+    def execute(self) -> list[object]:
         lock = self.backend._lock if self.transaction else RLock()
         with lock:
             results = [
@@ -171,7 +171,7 @@ class DisabledCacheBackend(InMemoryCacheBackend):
     def get(self, key: str) -> None:
         return None
 
-    def set(self, key: str, value: Any, **kwargs: Any) -> bool:
+    def set(self, key: str, value: str, **kwargs: object) -> bool:
         return False
 
     def delete(self, *keys: str) -> int:
diff --git a/policyengine_api/runtime_cache/reform_impacts.py b/policyengine_api/runtime_cache/reform_impacts.py
index 0076c8633..e21bdfe2b 100644
--- a/policyengine_api/runtime_cache/reform_impacts.py
+++ b/policyengine_api/runtime_cache/reform_impacts.py
@@ -85,7 +85,7 @@ def _datetime_to_wire(value: datetime | None) -> str | None:
     return value.astimezone(UTC).isoformat()
 
 
-def _datetime_from_wire(value: Any) -> datetime | None:
+def _datetime_from_wire(value: object) -> datetime | None:
     if value is None:
         return None
     if not isinstance(value, str):
@@ -103,7 +103,7 @@ def _impact_to_wire(impact: CachedReformImpact) -> dict[str, Any]:
     return values
 
 
-def _impact_from_wire(payload: Any) -> CachedReformImpact | None:
+def _impact_from_wire(payload: object) -> CachedReformImpact | None:
     if not isinstance(payload, dict):
         return None
     try:
@@ -532,7 +532,7 @@ def update(
         execution_id: str,
         *,
         clear_execution_id: bool = False,
-        **changes: Any,
+        **changes: Any,  # noqa: ANN401  # forwarded unchanged to dataclasses.replace
     ) -> CachedReformImpact | None:
         impact = self.get_by_execution_id(execution_id)
         if impact is None:
diff --git a/policyengine_api/worker_spm.py b/policyengine_api/worker_spm.py
index 2d9977d90..498b5c5af 100644
--- a/policyengine_api/worker_spm.py
+++ b/policyengine_api/worker_spm.py
@@ -1,5 +1,9 @@
 """Canonical economy execution requires the selected worker's certified capability."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
 from policyengine_api.spm import (
     SPMComparisonProvenance,
     SPMSelection,
@@ -10,14 +14,30 @@
     validate_spm_calculation_provenance,
 )
 
+if TYPE_CHECKING:
+    import httpx
+
+
+class SPMCapabilitySource(Protocol):
+    """A worker registry that reports its certified canonical SPM capability."""
+
+    def get_spm_capability(
+        self,
+        country: str,
+        model_version: str | None,
+        /,
+        *,
+        policyengine_version: str | None,
+    ) -> object: ...
+
 
 def validate_worker_spm(
     country_id: str,
     selection: object = None,
     *,
-    gateway=None,
-    policyengine_version=None,
-    model_version=None,
+    gateway: SPMCapabilitySource | None = None,
+    policyengine_version: str | None = None,
+    model_version: str | None = None,
 ) -> dict | None:
     country_id = country_id.lower() if isinstance(country_id, str) else country_id
     resolved = normalize_spm_selection(country_id, selection)
@@ -115,7 +135,7 @@ def validate_worker_result(
         ) from exc
 
 
-def raise_worker_spm_error(response) -> None:
+def raise_worker_spm_error(response: httpx.Response) -> None:
     """Preserve public typed input errors across asynchronous HTTP polling."""
     if response.status_code not in (400, 422):
         return
diff --git a/policyengine_api/worker_spm_release.py b/policyengine_api/worker_spm_release.py
index 79115b8c1..17d430325 100644
--- a/policyengine_api/worker_spm_release.py
+++ b/policyengine_api/worker_spm_release.py
@@ -21,16 +21,16 @@
 class RegistryCapability:
     """Expose one registry document through the request validator interface."""
 
-    def __init__(self, registry: dict):
+    def __init__(self, registry: dict) -> None:
         self.registry = registry
 
     def get_spm_capability(
         self,
-        country,
-        model_version,
+        country: str,
+        model_version: str | None,
         *,
-        policyengine_version,
-    ):
+        policyengine_version: str | None,
+    ) -> object:
         capabilities = self.registry.get("spm_capabilities")
         if not isinstance(capabilities, dict):
             return None

From 449394a449774485a7c29478ea767e90bfda0708 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:04:08 +0400
Subject: [PATCH 02/10] Annotate the economy, policy and household services and
 SPM helpers

Reform-impact service parameters take the cache's scalar types, with the
options hash optional on the write paths that receive it from
EconomicImpactSetupOptions. Policy JSON read from the v1 table is a
JSONValue. SPMSelection.serialize_selection stays without a return
annotation (noqa ANN201) because Pydantic would build the model's
serialization JSON schema from it and change the published OpenAPI.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 policyengine_api/country.py                   |  27 +++--
 .../services/economy_comparison.py            |   4 +-
 policyengine_api/services/economy_service.py  |  37 ++++---
 .../services/household_service.py             |   8 +-
 policyengine_api/services/policy_service.py   |   6 +-
 .../services/reform_impacts_service.py        | 100 +++++++++---------
 policyengine_api/spm.py                       |  30 ++++--
 7 files changed, 126 insertions(+), 86 deletions(-)

diff --git a/policyengine_api/country.py b/policyengine_api/country.py
index c1f754edf..93300b936 100644
--- a/policyengine_api/country.py
+++ b/policyengine_api/country.py
@@ -9,6 +9,7 @@
 from typing import Any
 
 import dpath
+import numpy as np
 import pandas as pd
 from policyengine_core.model_api import Enum, Reform
 from policyengine_core.parameters import (
@@ -19,7 +20,10 @@
     get_parameter,
 )
 from policyengine_core.periods import instant
+from policyengine_core.populations import Population
+from policyengine_core.simulations import Simulation
 from policyengine_core.taxbenefitsystems import TaxBenefitSystem
+from policyengine_core.variables import Variable
 
 from policyengine_api.constants import (
     get_bundle_default_dataset_option,
@@ -52,7 +56,7 @@ class PreparedCountryCalculation:
     spm_requested: bool
 
 
-def _serialize_float(value):
+def _serialize_float(value: object) -> float | str:
     serialized = float(str(value))
     if serialized == float("inf"):
         return "Infinity"
@@ -61,7 +65,12 @@ def _serialize_float(value):
     return serialized
 
 
-def _serialize_axis_result(result, variable, entity_index, count_entities):
+def _serialize_axis_result(
+    result: np.ndarray,
+    variable: Variable,
+    entity_index: int,
+    count_entities: int,
+) -> list[object]:
     if variable.value_type == Enum:
         result = result.decode_to_str()
 
@@ -79,7 +88,7 @@ def _record_calculation_failure(
     calculation_warnings: list[str],
     *,
     has_axes: bool,
-    population,
+    population: Population,
     entity_plural: str,
     entity_id: str,
     variable_name: str,
@@ -111,7 +120,7 @@ def _record_calculation_failure(
 
 
 class PolicyEngineCountry:
-    def __init__(self, country_package_name: str, country_id: str):
+    def __init__(self, country_package_name: str, country_id: str) -> None:
         self.country_package_name = country_package_name
         self.country_id = country_id
         self.country_package = importlib.import_module(country_package_name)
@@ -120,7 +129,7 @@ def __init__(self, country_package_name: str, country_id: str):
         )
         self.build_metadata()
 
-    def build_metadata(self):
+    def build_metadata(self) -> None:
         self.metadata = self._json_safe(
             dict(
                 variables=self.build_variables(),
@@ -144,7 +153,7 @@ def build_metadata(self):
             )
         )
 
-    def _json_safe(self, value):
+    def _json_safe(self, value: object) -> object:
         if isinstance(value, Path):
             return str(value)
         if isinstance(value, (date, datetime)):
@@ -567,7 +576,7 @@ def _create_simulation(
         household: dict,
         reform: dict | None,
         spm: dict | None = None,
-    ):
+    ) -> tuple[Simulation, TaxBenefitSystem]:
         selection = normalize_spm_selection(getattr(self, "country_id", ""), spm)
         normalized_reform = None
         if reform:
@@ -671,7 +680,7 @@ def modify_parameters(parameters: ParameterNode) -> ParameterNode:
         return parameters
 
     class reform(Reform):
-        def apply(self):
+        def apply(self) -> None:
             self.modify_parameters(modify_parameters)
 
     return reform
@@ -682,7 +691,7 @@ def get_requested_computations(
     *,
     include_provided_values: bool = False,
     variable_names: set[str] | None = None,
-):
+) -> list[tuple[str, str, str, str]]:
     requested_computations = dpath.util.search(
         household,
         "*/*/*/*",
diff --git a/policyengine_api/services/economy_comparison.py b/policyengine_api/services/economy_comparison.py
index a8be578f5..f0f992ee5 100644
--- a/policyengine_api/services/economy_comparison.py
+++ b/policyengine_api/services/economy_comparison.py
@@ -293,7 +293,9 @@ def poverty_impact(baseline: dict, reform: dict) -> dict:
     )
 
 
-def compute_income_change(baseline_values, reform_values):
+def compute_income_change(
+    baseline_values: np.ndarray, reform_values: np.ndarray
+) -> np.ndarray:
     """Percentage income change with a floor of 1 on the baseline
     to avoid division by zero for zero/negative incomes."""
     absolute_change = reform_values - baseline_values
diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py
index c6169f4cf..5979682c8 100644
--- a/policyengine_api/services/economy_service.py
+++ b/policyengine_api/services/economy_service.py
@@ -8,6 +8,7 @@
 import httpx
 from dotenv import load_dotenv
 from pydantic import BaseModel, Field
+from sqlalchemy.orm import Session, sessionmaker
 
 from policyengine_api.constants import (
     COUNTRY_PACKAGE_VERSIONS,
@@ -25,7 +26,13 @@
 from policyengine_api.data.places import validate_place_code
 from policyengine_api.data.v1_models import ReformImpact
 from policyengine_api.gcp_logging import logger
-from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint
+from policyengine_api.json_types import JSONValue
+from policyengine_api.libs.simulation_entrypoint import (
+    ModalBudgetWindowBatchExecution,
+    ModalSimulationExecution,
+    SimulationEntrypointClient,
+    simulation_entrypoint,
+)
 from policyengine_api.observability import (
     runtime as observability_runtime,
 )
@@ -44,6 +51,10 @@
     start_observability_id,
 )
 from policyengine_api.runtime_cache.core import CacheCoordinationError
+from policyengine_api.runtime_cache.reform_impacts import (
+    CachedReformImpact,
+    ReformImpactStartClaim,
+)
 from policyengine_api.services.budget_window_cache import BudgetWindowCache
 from policyengine_api.services.policy_service import PolicyService
 from policyengine_api.services.reform_impacts_service import (
@@ -262,11 +273,11 @@ class EconomyService:
     def __init__(
         self,
         *,
-        primary_session_factory=None,
+        primary_session_factory: sessionmaker[Session] | None = None,
         policy_service_: PolicyService | None = None,
         reform_impacts_service_: ReformImpactsService | None = None,
         budget_window_cache_: BudgetWindowCache | None = None,
-        simulation_entrypoint_=None,
+        simulation_entrypoint_: SimulationEntrypointClient | None = None,
     ) -> None:
         self._primary_session_factory = primary_session_factory
         self._injected_policy_service = policy_service_
@@ -291,7 +302,7 @@ def _budget_window_cache(self) -> BudgetWindowCache:
         return self._injected_budget_window_cache or budget_window_cache
 
     @property
-    def _simulation_gateway(self):
+    def _simulation_gateway(self) -> SimulationEntrypointClient:
         return self._injected_simulation_entrypoint or simulation_entrypoint
 
     @observability_runtime.span(ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_LOAD_POLICIES))
@@ -315,7 +326,7 @@ def _get_policy_jsons(
         )
 
     @staticmethod
-    def _parse_json_object(value: dict[str, Any] | str) -> dict[str, Any]:
+    def _parse_json_object(value: JSONValue) -> dict[str, Any]:
         """Accept ORM-decoded objects and legacy JSON text at a read boundary."""
 
         parsed = json.loads(value) if isinstance(value, str) else value
@@ -673,7 +684,7 @@ def _start_budget_window_batch(
         start_year: str,
         window_size: int,
         max_parallel: int,
-    ):
+    ) -> ModalBudgetWindowBatchExecution:
         sim_params = self._build_budget_window_batch_payload(
             setup_options=setup_options,
             start_year=start_year,
@@ -1106,7 +1117,7 @@ def _claim_reform_impact_start(
     def _get_reform_impact_start_claim(
         self,
         setup_options: EconomicImpactSetupOptions,
-    ):
+    ) -> ReformImpactStartClaim | None:
         return self._reform_impacts.get_reform_impact_start_claim(
             **self._reform_impact_start_claim_scope_arguments(setup_options)
         )
@@ -1149,7 +1160,7 @@ def _get_previous_impacts(
         time_period: str,
         options_hash: str,
         api_version: str,
-    ):
+    ) -> list[CachedReformImpact]:
         """
         Fetch any previous simulation runs for the given policy reform.
         """
@@ -1218,7 +1229,7 @@ def _handle_execution_state(
         setup_options: EconomicImpactSetupOptions,
         execution_state: str,
         reform_impact: ReformImpact,
-        execution: Any | None = None,
+        execution: ModalSimulationExecution | None = None,
     ) -> EconomicImpactResult:
         """
         Handle the state of the execution and return the appropriate status and result.
@@ -1590,7 +1601,7 @@ def _with_policyengine_bundle(
         self,
         result: dict,
         setup_options: EconomicImpactSetupOptions,
-        execution: Any | None = None,
+        execution: ModalSimulationExecution | None = None,
     ) -> dict:
         result = result if isinstance(result, dict) else {}
         cached_resolved_app_name = result.get("resolved_app_name")
@@ -1728,7 +1739,7 @@ def _set_reform_impact_computing(
         self,
         setup_options: EconomicImpactSetupOptions,
         execution_id: str,
-    ):
+    ) -> None:
         """
         In the reform_impact table, set the status of the impact to "computing".
         """
@@ -1766,7 +1777,7 @@ def _set_reform_impact_complete(
         setup_options: EconomicImpactSetupOptions,
         reform_impact_json: dict[str, Any],
         execution_id: str,
-    ):
+    ) -> None:
         """
         In the reform_impact table, set the status of the impact to "ok" and store the reform impact JSON.
         """
@@ -1800,7 +1811,7 @@ def _set_reform_impact_error(
         message: str,
         execution_id: str,
         error_code: str | None = None,
-    ):
+    ) -> None:
         """
         In the reform_impact table, set the status of the impact to "error" and store the error message.
         """
diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py
index 958dba053..1e67ec09c 100644
--- a/policyengine_api/services/household_service.py
+++ b/policyengine_api/services/household_service.py
@@ -35,8 +35,14 @@
 
     from sqlalchemy.orm import Session, sessionmaker
 
+    from policyengine_api.spm import SPMSelection
 
-def household_storage_json(country_id: str, household_json: dict, spm=None) -> dict:
+
+def household_storage_json(
+    country_id: str,
+    household_json: dict,
+    spm: SPMSelection | dict | None = None,
+) -> dict:
     """Keep the measurement selection in the same atomic JSON/hash as inputs.
 
     Only a selection the caller actually sent is stored. Certification still runs
diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py
index 3c5662605..5112d9429 100644
--- a/policyengine_api/services/policy_service.py
+++ b/policyengine_api/services/policy_service.py
@@ -2,7 +2,7 @@
 
 import copy
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from sqlalchemy import select
 
@@ -17,6 +17,8 @@
 
     from sqlalchemy.orm import Session, sessionmaker
 
+    from policyengine_api.json_types import JSONValue
+
 
 @dataclass(frozen=True)
 class PolicySetResult:
@@ -83,7 +85,7 @@ def get_policy_json(
         self,
         country_id: str,
         policy_id: int,
-    ) -> Any | None:
+    ) -> JSONValue:
         policy = self.get_policy(country_id, policy_id)
         return None if policy is None else policy.policy_json
 
diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py
index a046f060e..1658ff655 100644
--- a/policyengine_api/services/reform_impacts_service.py
+++ b/policyengine_api/services/reform_impacts_service.py
@@ -34,14 +34,14 @@ def get_recent_reform_impacts(
 
     def get_all_reform_impacts(
         self,
-        country_id,
-        policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
-        options_hash,
-        api_version,
+        country_id: str,
+        policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
+        options_hash: str,
+        api_version: str,
     ) -> list[CachedReformImpact]:
         return self._cache.matching(
             country_id=country_id,
@@ -56,15 +56,15 @@ def get_all_reform_impacts(
 
     def get_all_reform_impacts_by_options_hash_prefix(
         self,
-        country_id,
-        policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
-        options_hash,
-        options_hash_prefix,
-        api_version,
+        country_id: str,
+        policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
+        options_hash: str,
+        options_hash_prefix: str,
+        api_version: str,
     ) -> list[CachedReformImpact]:
         return self._cache.matching(
             country_id=country_id,
@@ -172,18 +172,18 @@ def release_reform_impact_start(
 
     def set_reform_impact(
         self,
-        country_id,
-        policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
+        country_id: str,
+        policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
         options: dict[str, Any],
-        options_hash,
-        status,
-        api_version,
+        options_hash: str | None,
+        status: str,
+        api_version: str,
         reform_impact_json: dict[str, Any],
-        start_time,
+        start_time: datetime.datetime,
         execution_id: str,
         observability_id: str | None = None,
     ) -> CachedReformImpact:
@@ -214,13 +214,13 @@ def set_reform_impact(
 
     def delete_reform_impact(
         self,
-        country_id,
-        policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
-        options_hash,
+        country_id: str,
+        policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
+        options_hash: str,
     ) -> None:
         self._cache.delete_matching_computing(
             country_id=country_id,
@@ -234,14 +234,14 @@ def delete_reform_impact(
 
     def set_error_reform_impact(
         self,
-        country_id,
-        policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
-        options_hash,
-        message,
+        country_id: str,
+        policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
+        options_hash: str | None,
+        message: str,
         execution_id: str,
         error_code: str | None = None,
     ) -> CachedReformImpact | None:
@@ -265,15 +265,15 @@ def set_error_reform_impact(
 
     def set_complete_reform_impact(
         self,
-        country_id,
-        reform_policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
-        options_hash,
+        country_id: str,
+        reform_policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
+        options_hash: str | None,
         reform_impact_json: dict[str, Any],
-        execution_id,
+        execution_id: str,
     ) -> CachedReformImpact | None:
         del (
             country_id,
diff --git a/policyengine_api/spm.py b/policyengine_api/spm.py
index 0130931ed..2dda6d2f5 100644
--- a/policyengine_api/spm.py
+++ b/policyengine_api/spm.py
@@ -11,7 +11,7 @@
 from collections.abc import Mapping, Sequence
 from datetime import date
 from functools import lru_cache
-from typing import Literal
+from typing import TYPE_CHECKING, Any, Literal, Self
 
 from pydantic import (
     BaseModel,
@@ -23,12 +23,19 @@
     model_validator,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from policyengine_core.simulations import Simulation
+    from pydantic import SerializerFunctionWrapHandler
+    from spm_calculator.rolling_forecast import SPMForecast
+
 
 class SPMSelection(BaseModel):
     """Select a forecast scenario and an explicit household SPM geography."""
 
     @staticmethod
-    def _selection_schema(schema):
+    def _selection_schema(schema: dict[str, Any]) -> None:
         # Omitted fields inherit the certified bundle, which may select a
         # different geography. Generated clients must preserve that omission.
         for field in schema.get("properties", {}).values():
@@ -72,7 +79,10 @@ def _selection_schema(schema):
     )
 
     @model_serializer(mode="wrap")
-    def serialize_selection(self, handler):
+    def serialize_selection(  # noqa: ANN201  # pydantic would build the output schema from it
+        self,
+        handler: SerializerFunctionWrapHandler,
+    ):
         """Keep omitted options distinct from explicit choices in nested JSON."""
         return {
             name: value
@@ -82,13 +92,13 @@ def serialize_selection(self, handler):
 
     @field_validator("as_of")
     @classmethod
-    def validate_as_of(cls, value):
+    def validate_as_of(cls, value: str | None) -> str | None:
         if value is not None and date.fromisoformat(value).isoformat() != value:
             raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)")
         return value
 
     @model_validator(mode="after")
-    def validate_location(self):
+    def validate_location(self) -> Self:
         if self.geography_kind == "metro":
             if not self.geography_id or not self.geography_id.strip():
                 raise ValueError("An SPM area selection requires geography_id")
@@ -314,7 +324,7 @@ def build_spm_comparison_provenance(
 class SPMValidationError(ValueError):
     """An SPM request/configuration error safe for the API's 4xx envelope."""
 
-    def __init__(self, code: str, message: str):
+    def __init__(self, code: str, message: str) -> None:
         self.code = code
         self.message = message
         super().__init__(message)
@@ -381,7 +391,7 @@ def spm_error_detail(error: BaseException) -> dict[str, str] | None:
     return None
 
 
-def simulation_supports_spm(simulation_type) -> bool:
+def simulation_supports_spm(simulation_type: Callable[..., object]) -> bool:
     """Old US constructors also accept **kwargs, so that alone is insufficient."""
     parameters = inspect.signature(simulation_type).parameters
     return "spm" in parameters or (
@@ -458,7 +468,7 @@ def _installed_country_implements_spm(country_id: str) -> bool:
 
 
 @lru_cache(maxsize=4)
-def _selected_forecast(expected_sha256: str):
+def _selected_forecast(expected_sha256: str) -> SPMForecast:
     from spm_calculator.rolling_forecast import load_forecast
 
     return load_forecast(expected_sha256=expected_sha256)
@@ -574,7 +584,7 @@ def normalize_spm_selection(
     return resolved.model_dump(mode="json")
 
 
-def _validate_metro_selection(forecast, selection: SPMSelection) -> None:
+def _validate_metro_selection(forecast: SPMForecast, selection: SPMSelection) -> None:
     """A selected area must exist in the artifact, without guessing a year.
 
     Area definitions may change between years. Accept an area available in any
@@ -617,7 +627,7 @@ def spm_metadata(country_id: str) -> dict:
     }
 
 
-def calculation_spm_receipt(simulation) -> dict:
+def calculation_spm_receipt(simulation: Simulation) -> dict:
     """Reduce a country's existing receipt; never request calculations here."""
     if not hasattr(simulation, "spm_config"):
         return {}

From d0729c898955ab9cd01e645e5a0c2a0a605521f6 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:04:08 +0400
Subject: [PATCH 03/10] Annotate v2 validators, query parameters and metadata
 reads

Before-validators take and return object. The parameter-tree SQL helpers
take mapped columns and return typed column expressions. Three helpers
keep Any with a reason: the JSON validators and canonical encoders
dispatch on exact type() checks mypy cannot narrow, and read_rows passes
core SQLAlchemy selects that Session.exec's overloads do not declare.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../data/v2/catalog/extraction.py             |  9 ++++++---
 .../v2/households/document_models.py          | 11 +++++++---
 policyengine_api/observability/identifiers.py |  3 +--
 policyengine_api/query_parameters.py          |  6 +++---
 .../services/v2/households/transformations.py |  4 +++-
 .../services/v2/households/validators.py      |  4 ++--
 .../v2/metadata/database_connectors/reads.py  |  5 ++++-
 .../reads_parameter_tree.py                   | 20 ++++++++++++++-----
 .../services/v2/policies/transformations.py   |  4 +++-
 .../services/v2/policies/validators.py        |  4 ++--
 10 files changed, 47 insertions(+), 23 deletions(-)

diff --git a/policyengine_api/data/v2/catalog/extraction.py b/policyengine_api/data/v2/catalog/extraction.py
index 0f04a4c4c..421d4633b 100644
--- a/policyengine_api/data/v2/catalog/extraction.py
+++ b/policyengine_api/data/v2/catalog/extraction.py
@@ -10,7 +10,7 @@
 from datetime import UTC, date, datetime, timedelta
 from enum import Enum
 from importlib import metadata as importlib_metadata
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from urllib.parse import urlsplit
 from uuid import NAMESPACE_URL, UUID, uuid5
 
@@ -28,6 +28,9 @@
     VariableRecord,
 )
 
+if TYPE_CHECKING:
+    from policyengine_api.json_types import JSONValue
+
 SUPPORTED_COUNTRIES = ("us", "uk")
 REQUIRED_DEPENDENCIES = (
     "policyengine-core",
@@ -137,7 +140,7 @@ def _type_name(value: object) -> str | None:
     return name
 
 
-def normalize_json_value(value: object) -> Any:
+def normalize_json_value(value: object) -> JSONValue:
     """Convert supported public values into strict JSON-compatible values."""
 
     if value is None or isinstance(value, (str, bool, int)):
@@ -153,7 +156,7 @@ def normalize_json_value(value: object) -> Any:
     if isinstance(value, (date, datetime)):
         return value.isoformat()
     if isinstance(value, Mapping):
-        normalized: dict[str, Any] = {}
+        normalized: dict[str, JSONValue] = {}
         for key, nested in value.items():
             if not isinstance(key, str):
                 raise CatalogExtractionError("JSON object keys must be strings")
diff --git a/policyengine_api/fastapi_routes/v2/households/document_models.py b/policyengine_api/fastapi_routes/v2/households/document_models.py
index 25a2d4dbc..d2b702b52 100644
--- a/policyengine_api/fastapi_routes/v2/households/document_models.py
+++ b/policyengine_api/fastapi_routes/v2/households/document_models.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Annotated, Any, TypeAlias
+from typing import TYPE_CHECKING, Annotated, Any, TypeAlias
 
 from pydantic import (
     BaseModel,
@@ -17,6 +17,9 @@
 
 from policyengine_api.spm import SPMSelection
 
+if TYPE_CHECKING:
+    from pydantic import SerializerFunctionWrapHandler
+
 DocumentIdentifier = Annotated[
     str,
     StringConstraints(strip_whitespace=False, min_length=1, max_length=255),
@@ -60,13 +63,15 @@ class USHouseholdDocument(StrictHouseholdDocumentModel):
 
     @field_validator("spm", mode="before")
     @classmethod
-    def require_spm_object(cls, value: Any) -> Any:
+    def require_spm_object(cls, value: object) -> object:
         if value is None:
             raise ValueError("spm must be an object when supplied")
         return value
 
     @model_serializer(mode="wrap")
-    def serialize_document(self, handler: Any) -> dict[str, Any]:
+    def serialize_document(
+        self, handler: SerializerFunctionWrapHandler
+    ) -> dict[str, Any]:
         document: dict[str, Any] = handler(self)
         if self.spm is None:
             document.pop("spm", None)
diff --git a/policyengine_api/observability/identifiers.py b/policyengine_api/observability/identifiers.py
index e836fc887..57be0c593 100644
--- a/policyengine_api/observability/identifiers.py
+++ b/policyengine_api/observability/identifiers.py
@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from typing import Any
 from uuid import UUID, uuid4
 
 OBSERVABILITY_ID_HEADER = "X-PolicyEngine-Observability-Id"
@@ -14,7 +13,7 @@ def generate_observability_id() -> str:
     return str(uuid4())
 
 
-def normalize_observability_id(value: Any) -> str | None:
+def normalize_observability_id(value: object) -> str | None:
     """Return a canonical UUID string, or ``None`` for malformed input."""
 
     if not isinstance(value, str):
diff --git a/policyengine_api/query_parameters.py b/policyengine_api/query_parameters.py
index 3b0d854cb..c11946886 100644
--- a/policyengine_api/query_parameters.py
+++ b/policyengine_api/query_parameters.py
@@ -43,7 +43,7 @@
 SUPPORTED_COUNTRY_IDS = frozenset({"us", "uk"})
 
 
-def normalize_country_id(value: Any) -> Any:
+def normalize_country_id(value: object) -> object:
     """Lowercase a textual country ID before its supported-value check."""
 
     return value.lower() if isinstance(value, str) else value
@@ -117,7 +117,7 @@ def _unique_json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
     return result
 
 
-def parse_spm_query_value(value: Any) -> Any:
+def parse_spm_query_value(value: object) -> dict[str, object] | SPMSelection:
     """Decode a single JSON object, preserving omissions and rejecting duplicates."""
     if isinstance(value, str):
         value = json.loads(value, object_pairs_hook=_unique_json_object)
@@ -126,7 +126,7 @@ def parse_spm_query_value(value: Any) -> Any:
     return value
 
 
-def parse_integer_query_value(value: Any) -> Any:
+def parse_integer_query_value(value: object) -> object:
     """Preserve integer-string parsing; decimal spellings such as 2.0 are invalid."""
     return int(value) if isinstance(value, str) else value
 
diff --git a/policyengine_api/services/v2/households/transformations.py b/policyengine_api/services/v2/households/transformations.py
index 25b8acca3..def489d9d 100644
--- a/policyengine_api/services/v2/households/transformations.py
+++ b/policyengine_api/services/v2/households/transformations.py
@@ -49,7 +49,9 @@ def _canonical_number(value: float) -> str:
     return rendered
 
 
-def _canonical_json(value: Any) -> str:
+def _canonical_json(
+    value: Any,  # noqa: ANN401  # exact type() dispatch that mypy cannot narrow
+) -> str:
     if value is None:
         return "null"
     if type(value) is bool:
diff --git a/policyengine_api/services/v2/households/validators.py b/policyengine_api/services/v2/households/validators.py
index ec031e71b..4f3e80fec 100644
--- a/policyengine_api/services/v2/households/validators.py
+++ b/policyengine_api/services/v2/households/validators.py
@@ -46,11 +46,11 @@ class LegacyHouseholdTranslationError(ValueError):
 
 
 def require_json_value(
-    value: Any,
+    value: Any,  # noqa: ANN401  # exact type() dispatch that mypy cannot narrow
     *,
     depth: int = 0,
     containers: frozenset[int] = frozenset(),
-) -> Any:
+) -> object:
     """Reject values that are not standards-compliant JSON."""
 
     if depth > MAXIMUM_JSON_NESTING:
diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads.py b/policyengine_api/services/v2/metadata/database_connectors/reads.py
index bc202807f..48690c061 100644
--- a/policyengine_api/services/v2/metadata/database_connectors/reads.py
+++ b/policyengine_api/services/v2/metadata/database_connectors/reads.py
@@ -31,7 +31,10 @@ def read_metadata_catalog(
     )
 
 
-def read_rows(session: Session, statement: Any) -> list[Any]:
+def read_rows(
+    session: Session,
+    statement: Any,  # noqa: ANN401  # Session.exec overloads omit core SQLAlchemy selects
+) -> list[Any]:
     try:
         return list(session.exec(statement).all())
     except SQLAlchemyError as error:
diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py b/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py
index 880df9f75..230793152 100644
--- a/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py
+++ b/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py
@@ -13,12 +13,16 @@
 if TYPE_CHECKING:
     from uuid import UUID
 
+    from sqlalchemy.orm import Mapped
+
 
 def _escaped_like(value: str) -> str:
     return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
 
 
-def _path_segment(remainder: Any, dialect: str) -> Any:
+def _path_segment(
+    remainder: sa.ColumnElement[str], dialect: str
+) -> sa.ColumnElement[str]:
     dot_position = (
         sa.func.instr(remainder, ".")
         if dialect == "sqlite"
@@ -30,17 +34,23 @@ def _path_segment(remainder: Any, dialect: str) -> Any:
     )
 
 
-def _child_path(column: Any, prefix: str, dialect: str) -> Any:
+def _child_path(
+    column: Mapped[str], prefix: str, dialect: str
+) -> sa.ColumnElement[str]:
     remainder = sa.func.substr(column, len(prefix) + 1)
     return sa.literal(prefix) + _path_segment(remainder, dialect)
 
 
-def _direct_child_path(column: Any, parent_path: Any, dialect: str) -> Any:
+def _direct_child_path(
+    column: Mapped[str], parent_path: sa.ColumnElement[str], dialect: str
+) -> sa.ColumnElement[str]:
     remainder = sa.func.substr(column, sa.func.length(parent_path) + 2)
     return parent_path + "." + _path_segment(remainder, dialect)
 
 
-def _has_path_prefix(column: Any, parent_path: Any) -> Any:
+def _has_path_prefix(
+    column: Mapped[str], parent_path: sa.ColumnElement[str]
+) -> sa.ColumnElement[bool]:
     return (
         sa.func.substr(column, 1, sa.func.length(parent_path) + 1) == parent_path + "."
     )
@@ -53,7 +63,7 @@ def _parameter_children_statement(
     dialect: str,
     offset: int,
     limit: int,
-) -> Any:
+) -> sa.Select[Any]:
     prefix = f"{parent_path}." if parent_path else ""
     escaped_prefix = _escaped_like(prefix)
     node_name = col(ParameterNode.name)
diff --git a/policyengine_api/services/v2/policies/transformations.py b/policyengine_api/services/v2/policies/transformations.py
index 839c7466d..f80e04390 100644
--- a/policyengine_api/services/v2/policies/transformations.py
+++ b/policyengine_api/services/v2/policies/transformations.py
@@ -52,7 +52,9 @@ def _canonical_number(value: float) -> str:
     return rendered
 
 
-def _canonical_json(value: Any) -> str:
+def _canonical_json(
+    value: Any,  # noqa: ANN401  # exact type() dispatch that mypy cannot narrow
+) -> str:
     if value is None:
         return "null"
     if type(value) is bool:
diff --git a/policyengine_api/services/v2/policies/validators.py b/policyengine_api/services/v2/policies/validators.py
index bc3afa613..0ceabb046 100644
--- a/policyengine_api/services/v2/policies/validators.py
+++ b/policyengine_api/services/v2/policies/validators.py
@@ -46,11 +46,11 @@ class LegacyPolicyTranslationError(ValueError):
 
 
 def require_json_value(
-    value: Any,
+    value: Any,  # noqa: ANN401  # exact type() dispatch that mypy cannot narrow
     *,
     depth: int = 0,
     containers: frozenset[int] = frozenset(),
-) -> Any:
+) -> object:
     """Reject values that cannot be represented as standards-compliant JSON."""
 
     if depth > MAXIMUM_JSON_NESTING:

From b8d2224bec31d1c7c7b22df75787bc2067a4d5ce Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:04:08 +0400
Subject: [PATCH 04/10] Annotate the Flask routes, ASGI shell and shared
 utilities

View decorators are generic over the wrapped view's result, so a view
returning dict | Response keeps that type plus the error Response.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 policyengine_api/api.py                       |  2 +-
 policyengine_api/asgi_factory.py              | 15 ++++++----
 policyengine_api/constants.py                 |  2 +-
 policyengine_api/data/orm.py                  |  6 +++-
 policyengine_api/migration_logging.py         |  4 +--
 policyengine_api/response_factory.py          |  3 +-
 policyengine_api/routes/economy_routes.py     |  6 ++--
 policyengine_api/routes/error_routes.py       | 10 +++----
 policyengine_api/routes/household_routes.py   | 30 ++++++++++++++-----
 policyengine_api/specification.py             | 10 ++++---
 policyengine_api/utils/cache_utils.py         |  2 +-
 policyengine_api/utils/deprecated_inputs.py   |  2 +-
 policyengine_api/utils/input_validation.py    |  4 +--
 policyengine_api/utils/json.py                | 12 ++++++--
 .../payload_validators/validate_country.py    | 20 ++++++++++---
 .../validate_household_payload.py             | 12 +++++++-
 policyengine_api/utils/singleton.py           |  2 +-
 17 files changed, 99 insertions(+), 43 deletions(-)

diff --git a/policyengine_api/api.py b/policyengine_api/api.py
index 881f104df..81fc92106 100644
--- a/policyengine_api/api.py
+++ b/policyengine_api/api.py
@@ -9,7 +9,7 @@
 start_time = time.time()
 
 
-def log_timing(message):
+def log_timing(message: str) -> None:
     elapsed = time.time() - start_time
     print(f"[{elapsed:.2f}s] {message}", file=sys.stderr, flush=True)
 
diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py
index 37a1da94c..f6d5afb3b 100644
--- a/policyengine_api/asgi_factory.py
+++ b/policyengine_api/asgi_factory.py
@@ -51,9 +51,11 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
 
+    from a2wsgi.wsgi_typing import WSGIApp
     from policyengine_observability import ObservabilityRuntime
+    from starlette.routing import BaseRoute
     from starlette.types import ASGIApp
 
 
@@ -71,7 +73,7 @@ def _apply_observability_id_header(
     response.headers[OBSERVABILITY_ID_HEADER] = observability_id
 
 
-def _matched_route_template(routes, scope: dict) -> str | None:
+def _matched_route_template(routes: Iterable[BaseRoute], scope: dict) -> str | None:
     for route in routes:
         match, _ = route.matches(scope)
         if match is Match.FULL:
@@ -100,7 +102,7 @@ def _native_route_template(app: FastAPI, scope: dict) -> str | None:
 
 
 def create_asgi_app(
-    wsgi_app,
+    wsgi_app: WSGIApp,
     *,
     route_settings: RouteImplementationSettings | None = None,
     dependencies: NativeRouteDependencies | None = None,
@@ -116,7 +118,7 @@ def create_asgi_app(
     request_runtime = observability_runtime or get_runtime()
 
     @asynccontextmanager
-    async def lifespan(_app: FastAPI):
+    async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
         try:
             yield
         finally:
@@ -172,7 +174,10 @@ async def oversized_v2_request(
         return v2_error_response(413, str(error))
 
     @app.middleware("http")
-    async def add_request_context_and_migration_logging(request, call_next):
+    async def add_request_context_and_migration_logging(
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
         started_at = time.time()
         request_id = request.headers.get(REQUEST_ID_HEADER) or generate_request_id()
         incoming_observability_id = normalize_observability_id(
diff --git a/policyengine_api/constants.py b/policyengine_api/constants.py
index 7056574cb..6aeef3fda 100644
--- a/policyengine_api/constants.py
+++ b/policyengine_api/constants.py
@@ -47,7 +47,7 @@ def _resolve_distribution_version(
     return "0.0.0"
 
 
-def get_py_manifest():
+def get_py_manifest() -> Path:
     return distribution("policyengine").locate_file(
         "policyengine/data/bundle/manifest.json"
     )
diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py
index 6c0501e29..1a1ffaa67 100644
--- a/policyengine_api/data/orm.py
+++ b/policyengine_api/data/orm.py
@@ -4,12 +4,16 @@
 
 import atexit
 import os
+from typing import TYPE_CHECKING
 
 from dotenv import load_dotenv
 from google.cloud.sql.connector import Connector, IPTypes
 from sqlalchemy import Engine, create_engine
 from sqlalchemy.orm import Session, sessionmaker
 
+if TYPE_CHECKING:
+    from sqlalchemy.engine.interfaces import DBAPIConnection
+
 load_dotenv()
 
 DEFAULT_REMOTE_DB_USER = "policyengine"
@@ -49,7 +53,7 @@ def _build_remote_engine() -> Engine:
     )
     password = os.environ["POLICYENGINE_DB_PASSWORD"]
 
-    def get_connection():
+    def get_connection() -> DBAPIConnection:
         return connector.connect(
             instance_connection_string=config["instance_connection_name"],
             driver="pymysql",
diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py
index 9123a53a4..c521e77a8 100644
--- a/policyengine_api/migration_logging.py
+++ b/policyengine_api/migration_logging.py
@@ -82,7 +82,7 @@ def register_migration_request_logging(
     """Register request IDs and migration logging for Flask."""
 
     @app.before_request
-    def set_request_migration_context():
+    def set_request_migration_context() -> None:
         flask.g.request_started_at = time.time()
         try:
             captured = runtime.capture_context() if runtime is not None else {}
@@ -102,7 +102,7 @@ def set_request_migration_context():
                 pass
 
     @app.after_request
-    def log_request_migration_context(response):
+    def log_request_migration_context(response: flask.Response) -> flask.Response:
         request_id = getattr(flask.g, "request_id", None)
         if request_id is not None:
             response.headers[REQUEST_ID_HEADER] = request_id
diff --git a/policyengine_api/response_factory.py b/policyengine_api/response_factory.py
index 84cdd4240..08b17c0db 100644
--- a/policyengine_api/response_factory.py
+++ b/policyengine_api/response_factory.py
@@ -1,7 +1,6 @@
 """Factories for consistent Flask response construction."""
 
 import json
-from typing import Any
 
 from flask import Response
 
@@ -12,7 +11,7 @@ def _make_error_response(
     *,
     include_status: bool = True,
     mimetype: str | None = "application/json",
-    **payload_fields: Any,
+    **payload_fields: object,
 ) -> Response:
     """Build a JSON error response while preserving legacy payload variants."""
     payload = {"message": str(message), **payload_fields}
diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py
index 60837fa7d..cf493aac7 100644
--- a/policyengine_api/routes/economy_routes.py
+++ b/policyengine_api/routes/economy_routes.py
@@ -70,7 +70,9 @@ def _bad_request_response(error: str | ValueError) -> Response:
     methods=["GET"],
 )
 @validate_country
-def get_economic_impact(country_id: str, policy_id: int, baseline_policy_id: int):
+def get_economic_impact(
+    country_id: str, policy_id: int, baseline_policy_id: int
+) -> Response:
     policy_id = int(policy_id or get_current_law_policy_id(country_id))
     baseline_policy_id = int(
         baseline_policy_id or get_current_law_policy_id(country_id)
@@ -119,7 +121,7 @@ def get_economic_impact(country_id: str, policy_id: int, baseline_policy_id: int
 @validate_country
 def get_budget_window_economic_impact(
     country_id: str, policy_id: int, baseline_policy_id: int
-):
+) -> Response:
     policy_id = int(policy_id or get_current_law_policy_id(country_id))
     baseline_policy_id = int(
         baseline_policy_id or get_current_law_policy_id(country_id)
diff --git a/policyengine_api/routes/error_routes.py b/policyengine_api/routes/error_routes.py
index a4801bb09..b5f7ddb5f 100644
--- a/policyengine_api/routes/error_routes.py
+++ b/policyengine_api/routes/error_routes.py
@@ -9,31 +9,31 @@
 
 
 @error_bp.app_errorhandler(404)
-def response_404(error) -> Response:
+def response_404(error: HTTPException) -> Response:
     """Specific handler for 404 Not Found errors"""
     return _make_error_response(error, 404, result=None)
 
 
 @error_bp.app_errorhandler(400)
-def response_400(error) -> Response:
+def response_400(error: HTTPException) -> Response:
     """Specific handler for 400 Bad Request errors"""
     return _make_error_response(error, 400, result=None)
 
 
 @error_bp.app_errorhandler(401)
-def response_401(error) -> Response:
+def response_401(error: HTTPException) -> Response:
     """Specific handler for 401 Unauthorized errors"""
     return _make_error_response(error, 401, result=None)
 
 
 @error_bp.app_errorhandler(403)
-def response_403(error) -> Response:
+def response_403(error: HTTPException) -> Response:
     """Specific handler for 403 Forbidden errors"""
     return _make_error_response(error, 403, result=None)
 
 
 @error_bp.app_errorhandler(500)
-def response_500(error) -> Response:
+def response_500(error: HTTPException) -> Response:
     """Specific handler for 500 Internal Server errors"""
     return _make_error_response(error, 500, result=None)
 
diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py
index 13c8888b4..7fe178e6f 100644
--- a/policyengine_api/routes/household_routes.py
+++ b/policyengine_api/routes/household_routes.py
@@ -1,7 +1,9 @@
 import json
 import logging
 import time
+from collections.abc import Callable
 from functools import wraps
+from typing import TypeVar
 
 from flask import Blueprint, Response, g, request
 from policyengine_core.errors import SituationParsingError
@@ -22,6 +24,7 @@
 )
 from policyengine_api.response_factory import _make_error_response
 from policyengine_api.services.household_calculation_service import (
+    HouseholdCalculationResult,
     HouseholdCalculationService,
     HouseholdNotFoundError,
     InvalidHouseholdInputsError,
@@ -47,6 +50,8 @@
     validate_household_payload,
 )
 
+_ViewResult = TypeVar("_ViewResult")
+
 household_bp = Blueprint("household", __name__)
 household_service = HouseholdService()
 household_calculation_service = HouseholdCalculationService()
@@ -133,7 +138,7 @@ def _spm_error_response(error: Exception) -> Response | None:
     return None
 
 
-def _calculation_response(calculation) -> dict:
+def _calculation_response(calculation: HouseholdCalculationResult) -> dict:
     result = dict(status="ok", message=None, result=calculation.household)
     if calculation.warnings:
         result["warnings"] = list(calculation.warnings)
@@ -142,7 +147,7 @@ def _calculation_response(calculation) -> dict:
     return result
 
 
-def _requested_spm(country_id: str, payload: dict):
+def _requested_spm(country_id: str, payload: dict) -> object:
     """Read a chosen selection, treating an explicit null as a malformed one.
 
     A v2 document and a simulation record both refuse an explicit null rather
@@ -169,12 +174,21 @@ def _requested_spm(country_id: str, payload: dict):
     return selection
 
 
-def _validate_calculation_request(*, add_missing: bool):
+def _validate_calculation_request(
+    *, add_missing: bool
+) -> Callable[
+    [Callable[..., _ViewResult]],
+    Callable[..., _ViewResult | Response],
+]:
     """Validate calculation inputs before an HTTP cache can satisfy a request."""
 
-    def decorator(func):
+    def decorator(
+        func: Callable[..., _ViewResult],
+    ) -> Callable[..., _ViewResult | Response]:
         @wraps(func)
-        def wrapped(country_id, *args, **kwargs):
+        def wrapped(
+            country_id: str, *args: object, **kwargs: object
+        ) -> _ViewResult | Response:
             payload = request.get_json()
             if not isinstance(payload, dict):
                 raise BadRequest("Calculation payload must be a JSON object.")
@@ -225,7 +239,7 @@ def wrapped(country_id, *args, **kwargs):
     return decorator
 
 
-def _calculation_cache_key(*args, **kwargs):
+def _calculation_cache_key(*args: object, **kwargs: object) -> str:
     country_id = request.view_args["country_id"]
     return hash_object(
         {
@@ -354,7 +368,9 @@ def post_household(country_id: str) -> Response:
     methods=["GET"],
 )
 @validate_country
-def get_household_under_policy(country_id: str, household_id: str, policy_id: str):
+def get_household_under_policy(
+    country_id: str, household_id: str, policy_id: str
+) -> dict | Response:
     """Get a stored household's output under a stored policy."""
     try:
         get_v1_household_read_source()
diff --git a/policyengine_api/specification.py b/policyengine_api/specification.py
index 986988ed7..adc5d6631 100644
--- a/policyengine_api/specification.py
+++ b/policyengine_api/specification.py
@@ -22,12 +22,14 @@
 )
 
 if TYPE_CHECKING:
-    from policyengine_api.json_types import JSONObject
+    from pydantic import BaseModel
+
+    from policyengine_api.json_types import JSONObject, JSONValue
 
 DEFAULT_SPECIFICATION_PATH = Path(__file__).with_name("openapi_spec.yaml")
 
 
-def _openapi_30_schema(value):
+def _openapi_30_schema(value: JSONValue) -> JSONValue:
     """Convert Pydantic's optional JSON Schema fields to OpenAPI 3.0 nullable."""
     if isinstance(value, list):
         return [_openapi_30_schema(item) for item in value]
@@ -47,7 +49,7 @@ def _openapi_30_schema(value):
     return result
 
 
-def _economy_query_parameters(model):
+def _economy_query_parameters(model: type[BaseModel]) -> list[dict[str, object]]:
     """Publish the same complete scalar query contract used by Flask."""
     schema = model.model_json_schema()
     result = []
@@ -70,7 +72,7 @@ def _economy_query_parameters(model):
     return result
 
 
-def _component_schema(model):
+def _component_schema(model: type[BaseModel]) -> JSONValue:
     """Generate one OpenAPI component and publish nested models separately."""
     schema = model.model_json_schema(
         ref_template="#/components/schemas/{model}",
diff --git a/policyengine_api/utils/cache_utils.py b/policyengine_api/utils/cache_utils.py
index c6da3812a..87d74b66c 100644
--- a/policyengine_api/utils/cache_utils.py
+++ b/policyengine_api/utils/cache_utils.py
@@ -7,7 +7,7 @@
 import flask
 
 
-def make_cache_key(*args, **kwargs):
+def make_cache_key(*args: object, **kwargs: object) -> str:
     # pylint: disable=unused-argument
     """make a hash to uniquely identify a cache entry.
     keep it fast, adding overhead to try to add some minor chance of a
diff --git a/policyengine_api/utils/deprecated_inputs.py b/policyengine_api/utils/deprecated_inputs.py
index 2b45b4573..6d6810c46 100644
--- a/policyengine_api/utils/deprecated_inputs.py
+++ b/policyengine_api/utils/deprecated_inputs.py
@@ -135,7 +135,7 @@ def _drop_deprecated_axes(
 
 
 def _is_deprecated_axis(
-    axis, location: str, warnings: list[DeprecatedVariableWarning]
+    axis: object, location: str, warnings: list[DeprecatedVariableWarning]
 ) -> bool:
     if not isinstance(axis, dict):
         return False
diff --git a/policyengine_api/utils/input_validation.py b/policyengine_api/utils/input_validation.py
index 68daa3c5f..a74934197 100644
--- a/policyengine_api/utils/input_validation.py
+++ b/policyengine_api/utils/input_validation.py
@@ -156,7 +156,7 @@ def format_unrecognized_inputs_message(errors: list[UnrecognizedInput]) -> str:
     )
 
 
-def _find_unrecognized_axes(axes, variables: dict) -> list[UnrecognizedInput]:
+def _find_unrecognized_axes(axes: object, variables: dict) -> list[UnrecognizedInput]:
     if not isinstance(axes, list):
         return []
 
@@ -183,7 +183,7 @@ def _find_unrecognized_axes(axes, variables: dict) -> list[UnrecognizedInput]:
 
 
 def _find_unrecognized_axis(
-    axis,
+    axis: object,
     path: str,
     variables: dict,
 ) -> list[UnrecognizedInput]:
diff --git a/policyengine_api/utils/json.py b/policyengine_api/utils/json.py
index 294e797db..9b617b95a 100644
--- a/policyengine_api/utils/json.py
+++ b/policyengine_api/utils/json.py
@@ -1,10 +1,16 @@
+from __future__ import annotations
+
 import base64
 import hashlib
+from typing import TYPE_CHECKING
 
 import numpy as np
 
+if TYPE_CHECKING:
+    from policyengine_api.json_types import JSONValue
+
 
-def make_hashable(o):
+def make_hashable(o: object) -> object:
     if isinstance(o, (tuple, list)):
         return tuple(make_hashable(e) for e in o)
 
@@ -17,13 +23,13 @@ def make_hashable(o):
     return o
 
 
-def hash_object(o):
+def hash_object(o: object) -> str:
     hasher = hashlib.sha256()
     hasher.update(repr(make_hashable(o)).encode())
     return base64.b64encode(hasher.digest()).decode()
 
 
-def get_safe_json(value):
+def get_safe_json(value: object) -> JSONValue:
     # This function is used when constructing metadata,
     # which will be consumed by JS, hence Infinity, not .inf
     if isinstance(value, (int, float)):
diff --git a/policyengine_api/utils/payload_validators/validate_country.py b/policyengine_api/utils/payload_validators/validate_country.py
index adc2204d0..9fc4130c1 100644
--- a/policyengine_api/utils/payload_validators/validate_country.py
+++ b/policyengine_api/utils/payload_validators/validate_country.py
@@ -1,6 +1,7 @@
-from functools import wraps
+from __future__ import annotations
 
-from flask import Response
+from functools import wraps
+from typing import TYPE_CHECKING, TypeVar
 
 from policyengine_api.country_validation import (
     InvalidCountryError,
@@ -8,8 +9,17 @@
 )
 from policyengine_api.response_factory import _make_error_response
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from flask import Response
+
+_ViewResult = TypeVar("_ViewResult")
+
 
-def validate_country(func):
+def validate_country(
+    func: Callable[..., _ViewResult],
+) -> Callable[..., _ViewResult | Response]:
     """Validate that a country ID is valid. If not, return a 400 response.
 
     Args:
@@ -20,7 +30,9 @@ def validate_country(func):
     """
 
     @wraps(func)
-    def validate_country_wrapper(country_id: str, *args, **kwargs) -> None | Response:
+    def validate_country_wrapper(
+        country_id: str, *args: object, **kwargs: object
+    ) -> _ViewResult | Response:
         print("Validating country")
         try:
             ensure_supported_country(country_id)
diff --git a/policyengine_api/utils/payload_validators/validate_household_payload.py b/policyengine_api/utils/payload_validators/validate_household_payload.py
index e3d281aad..19a80a672 100644
--- a/policyengine_api/utils/payload_validators/validate_household_payload.py
+++ b/policyengine_api/utils/payload_validators/validate_household_payload.py
@@ -1,4 +1,14 @@
-def validate_household_payload(payload):
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
+def validate_household_payload(
+    payload: Mapping[str, object],
+) -> tuple[bool, str | None]:
     """
     Validate the payload for a POST request to set a household's input data.
 
diff --git a/policyengine_api/utils/singleton.py b/policyengine_api/utils/singleton.py
index 9767a5c7e..ecba9af47 100644
--- a/policyengine_api/utils/singleton.py
+++ b/policyengine_api/utils/singleton.py
@@ -1,7 +1,7 @@
 class Singleton(type):
     _instances = {}
 
-    def __call__(cls, *args, **kwargs):
+    def __call__(cls, *args: object, **kwargs: object) -> object:
         if cls not in cls._instances:
             cls._instances[cls] = super().__call__(*args, **kwargs)
         return cls._instances[cls]

From e440be7f57bcc456642d9c724e10983272a13726 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:04:08 +0400
Subject: [PATCH 05/10] Annotate the rules tool, migration scripts and Alembic
 environments

scripts/rules stays standard-library only. Enforcer.get keeps Any for
the raw TOML value that callers read by its shape. scripts/rules changes
need mirroring in policyengine-sim-api.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 dashboard/experiments/gpt4_api_user.py | 10 +++++-----
 migrations/v1/env.py                   |  8 +++++++-
 migrations/v2/env.py                   | 18 ++++++++++++++++--
 scripts/rules/guard.py                 |  4 ++--
 scripts/rules/load.py                  | 21 +++++++++++++--------
 scripts/rules/model.py                 |  6 +++++-
 scripts/v1_database_migration.py       |  3 +--
 7 files changed, 49 insertions(+), 21 deletions(-)

diff --git a/dashboard/experiments/gpt4_api_user.py b/dashboard/experiments/gpt4_api_user.py
index 61c13bd0e..27b21f890 100644
--- a/dashboard/experiments/gpt4_api_user.py
+++ b/dashboard/experiments/gpt4_api_user.py
@@ -13,7 +13,7 @@
 
 
 @st.cache_data
-def get_embeddings_df():
+def get_embeddings_df() -> pd.DataFrame:
     embeddings_df = pd.read_csv("parameter_embeddings.csv.gz", compression="gzip")
     embeddings_df.parameter_embedding = embeddings_df.parameter_embedding.apply(
         lambda x: eval(x)
@@ -22,7 +22,7 @@ def get_embeddings_df():
 
 
 @st.cache_data
-def embed(prompt, engine="text-embedding-ada-002"):
+def embed(prompt: str, engine: str = "text-embedding-ada-002") -> list[float]:
     return get_embedding(prompt, engine=engine)
 
 
@@ -58,7 +58,7 @@ def embed(prompt, engine="text-embedding-ada-002"):
 # display in streamlit
 
 
-def write_prompt(question, relevant_parameters):
+def write_prompt(question: str, relevant_parameters: object) -> str:
     return f"""Question: {question}
     Here's some metadata about the parameters that are relevant to your question:
     Relevant parameters: {relevant_parameters}
@@ -78,7 +78,7 @@ def write_prompt(question, relevant_parameters):
 
 
 @st.cache_data
-def get_gpt4_response(prompt):
+def get_gpt4_response(prompt: str) -> str:
     return openai.ChatCompletion.create(
         model="gpt-4",
         messages=[{"role": "user", "content": prompt}],
@@ -93,7 +93,7 @@ def get_gpt4_response(prompt):
 
 
 @st.cache_data
-def get_policyengine_impact(reform):
+def get_policyengine_impact(reform: object) -> object:
     payload = {"data": reform}
     url = "https://api.policyengine.org/uk/policy"
 
diff --git a/migrations/v1/env.py b/migrations/v1/env.py
index 1669dc5ed..ab4b6983e 100644
--- a/migrations/v1/env.py
+++ b/migrations/v1/env.py
@@ -1,7 +1,10 @@
 """Alembic environment for the API v1 schema."""
 
+from __future__ import annotations
+
 import os
 from logging.config import fileConfig
+from typing import TYPE_CHECKING
 
 from alembic import context
 from sqlalchemy import engine_from_config, pool
@@ -9,6 +12,9 @@
 
 from policyengine_api.data.v1_models import V1Base
 
+if TYPE_CHECKING:
+    from sqlalchemy import Connection
+
 config = context.config
 if config.config_file_name is not None:
     fileConfig(config.config_file_name)
@@ -59,7 +65,7 @@ def run_migrations_online() -> None:
         _run_migrations_with_connection(connection)
 
 
-def _run_migrations_with_connection(connection) -> None:
+def _run_migrations_with_connection(connection: Connection) -> None:
     context.configure(
         connection=connection,
         target_metadata=target_metadata,
diff --git a/migrations/v2/env.py b/migrations/v2/env.py
index 1ab23e003..17b8e7d6d 100644
--- a/migrations/v2/env.py
+++ b/migrations/v2/env.py
@@ -1,6 +1,9 @@
 """Alembic environment for the isolated API v2-alpha Postgres schema."""
 
+from __future__ import annotations
+
 from logging.config import fileConfig
+from typing import TYPE_CHECKING
 
 from alembic import context
 from alembic.script import ScriptDirectory
@@ -13,6 +16,11 @@
 )
 from policyengine_api.data.v2.models import V2_METADATA
 
+if TYPE_CHECKING:
+    from alembic.runtime.environment import NameFilterType
+    from sqlalchemy import Connection
+    from sqlalchemy.sql.schema import SchemaItem
+
 config = context.config
 if config.config_file_name is not None:
     fileConfig(config.config_file_name)
@@ -22,13 +30,19 @@
 script = ScriptDirectory.from_config(config)
 
 
-def _include_application_object(_object, name, type_, _reflected, _compare_to) -> bool:
+def _include_application_object(
+    _object: SchemaItem,
+    name: str | None,
+    type_: NameFilterType,
+    _reflected: bool,
+    _compare_to: SchemaItem | None,
+) -> bool:
     """Keep Alembic's own version table outside application drift."""
 
     return not (type_ == "table" and name == "alembic_version")
 
 
-def _configure(connection) -> None:
+def _configure(connection: Connection) -> None:
     qualify_v2_connection(connection, settings)
     context.configure(
         connection=connection,
diff --git a/scripts/rules/guard.py b/scripts/rules/guard.py
index 42d4b6a37..c68fce02b 100644
--- a/scripts/rules/guard.py
+++ b/scripts/rules/guard.py
@@ -16,7 +16,7 @@
 )
 from scripts.rules.discover import discover_hidden
 from scripts.rules.load import load_registry
-from scripts.rules.model import Violation
+from scripts.rules.model import Registry, Violation
 from scripts.rules.render import (
     AGENTS_MD,
     RULES_MD,
@@ -42,7 +42,7 @@ class GuardResult:
     notes: list[str] = field(default_factory=list)
 
 
-def _stale_outputs(registry, root: Path) -> list[Violation]:
+def _stale_outputs(registry: Registry, root: Path) -> list[Violation]:
     violations = []
     rules_md = root / RULES_MD
     if not rules_md.exists():
diff --git a/scripts/rules/load.py b/scripts/rules/load.py
index 267818647..83b455e2e 100644
--- a/scripts/rules/load.py
+++ b/scripts/rules/load.py
@@ -8,8 +8,9 @@
 
 import tomllib
 from datetime import date
+from enum import Enum
 from pathlib import Path
-from typing import Any
+from typing import Any, TypeVar
 
 from scripts.rules.model import (
     Agents,
@@ -30,6 +31,8 @@
     Violation,
 )
 
+_EnumT = TypeVar("_EnumT", bound=Enum)
+
 REGISTRY_PATH = Path("rules/registry.toml")
 
 _TOP_KEYS = {
@@ -127,7 +130,9 @@ def string_list(
         )
         return ()
 
-    def enum(self, rule_id: str | None, enum_type: Any, value: Any, where: str) -> Any:
+    def enum(
+        self, rule_id: str | None, enum_type: type[_EnumT], value: object, where: str
+    ) -> _EnumT | None:
         try:
             return enum_type(value)
         except ValueError:
@@ -135,7 +140,7 @@ def enum(self, rule_id: str | None, enum_type: Any, value: Any, where: str) -> A
             self.add(rule_id, f"{where} is {value!r}", f"use one of {allowed}")
             return None
 
-    def day(self, rule_id: str | None, value: Any, where: str) -> date | None:
+    def day(self, rule_id: str | None, value: object, where: str) -> date | None:
         if isinstance(value, date):
             return value
         if isinstance(value, str):
@@ -147,7 +152,7 @@ def day(self, rule_id: str | None, value: Any, where: str) -> date | None:
         return None
 
 
-def _ignored(c: _Collector, raw: Any, key: str, field: str) -> tuple[Ignored, ...]:
+def _ignored(c: _Collector, raw: object, key: str, field: str) -> tuple[Ignored, ...]:
     if raw is None:
         return ()
     items = []
@@ -168,7 +173,7 @@ def _ignored(c: _Collector, raw: Any, key: str, field: str) -> tuple[Ignored, ..
     return tuple(items)
 
 
-def _discovery(c: _Collector, raw: Any) -> Discovery:
+def _discovery(c: _Collector, raw: object) -> Discovery:
     if raw is None:
         return Discovery()
     if not isinstance(raw, dict):
@@ -189,7 +194,7 @@ def _discovery(c: _Collector, raw: Any) -> Discovery:
     )
 
 
-def _agents(c: _Collector, raw: Any) -> Agents | None:
+def _agents(c: _Collector, raw: object) -> Agents | None:
     if raw is None:
         return None
     if not isinstance(raw, dict):
@@ -207,7 +212,7 @@ def _agents(c: _Collector, raw: Any) -> Agents | None:
     )
 
 
-def _ci_scripts(c: _Collector, raw: Any) -> CiScripts | None:
+def _ci_scripts(c: _Collector, raw: object) -> CiScripts | None:
     if raw is None:
         return None
     if not isinstance(raw, dict):
@@ -225,7 +230,7 @@ def _ci_scripts(c: _Collector, raw: Any) -> CiScripts | None:
     )
 
 
-def _rule(c: _Collector, raw: Any, index: int) -> Rule | None:
+def _rule(c: _Collector, raw: object, index: int) -> Rule | None:
     if not isinstance(raw, dict):
         c.add(None, f"rule #{index + 1} must be a table", "write [[rule]]")
         return None
diff --git a/scripts/rules/model.py b/scripts/rules/model.py
index 619138b01..6d8f9c365 100644
--- a/scripts/rules/model.py
+++ b/scripts/rules/model.py
@@ -100,7 +100,11 @@ class Enforcer:
     kind: Kind
     fields: Mapping[str, Any] = field(default_factory=dict)
 
-    def get(self, key: str, default: Any = None) -> Any:
+    def get(
+        self,
+        key: str,
+        default: object = None,
+    ) -> Any:  # noqa: ANN401  # a raw TOML value that callers read by its shape
         return self.fields.get(key, default)
 
     @property
diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py
index 084792311..e57194a3e 100644
--- a/scripts/v1_database_migration.py
+++ b/scripts/v1_database_migration.py
@@ -5,7 +5,6 @@
 import argparse
 import os
 from enum import StrEnum
-from typing import Any
 
 from alembic import command
 from alembic.autogenerate import compare_metadata
@@ -70,7 +69,7 @@ def classify_database_state(
     return DatabaseState.PENDING
 
 
-def describe_metadata_difference(difference: Any) -> str:
+def describe_metadata_difference(difference: object) -> str:
     """Return a stable, data-free description of an Alembic metadata diff."""
 
     item = difference

From e0b220ed9f6582a5aa1914466090d322ebc7eb0b Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:04:26 +0400
Subject: [PATCH 06/10] Enable ruff group 3 (annotations)

Select ANN everywhere except tests, which need no annotations. The
annotation commits before this one fix every finding outside .github/.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyproject.toml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 3c54d85d2..f14d0d729 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -89,12 +89,17 @@ select = [
     "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W",
     # Group 2: typing-only imports.
     "TC",
+    # Group 3: annotations.
+    "ANN",
 ]
 
 [tool.ruff.lint.per-file-ignores]
 # Alembic revisions must stay exactly as Alembic generates them
 # (api.migration.alembic-autogenerated-only), so ruff does not lint them.
 "migrations/*/versions/*.py" = ["ALL"]
+# Tests may use pytest's bare assert, and need no docstrings, annotations or
+# named constants for magic numbers. Every other rule applies to tests.
+"**/tests/**" = ["ANN", "D1", "PLR2004", "S101"]
 
 [tool.mypy]
 python_version = "3.11"

From 15645b55d9ae284c67c75036830ee20c624a5d98 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:33:24 +0400
Subject: [PATCH 07/10] Correct annotations that whole-package mypy shows were
 wrong

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 policyengine_api/asgi_factory.py                | 11 +++++------
 policyengine_api/constants.py                   |  3 ++-
 policyengine_api/country.py                     | 16 ++++++++++------
 policyengine_api/routes/household_routes.py     |  4 +++-
 policyengine_api/services/economy_comparison.py | 12 ++++++------
 policyengine_api/services/economy_service.py    | 12 ++++++------
 policyengine_api/services/household_service.py  |  6 ++----
 policyengine_api/specification.py               | 10 +++++++---
 policyengine_api/spm.py                         |  2 +-
 policyengine_api/utils/cache_utils.py           |  3 ++-
 policyengine_api/utils/json.py                  |  6 ++++--
 11 files changed, 48 insertions(+), 37 deletions(-)

diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py
index f6d5afb3b..0c15fb9d6 100644
--- a/policyengine_api/asgi_factory.py
+++ b/policyengine_api/asgi_factory.py
@@ -4,7 +4,7 @@
 
 import time
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from a2wsgi import WSGIMiddleware
 from fastapi import FastAPI, Request
@@ -53,10 +53,9 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
 
-    from a2wsgi.wsgi_typing import WSGIApp
     from policyengine_observability import ObservabilityRuntime
     from starlette.routing import BaseRoute
-    from starlette.types import ASGIApp
+    from starlette.types import ASGIApp, Scope
 
 
 def _apply_request_id_header(
@@ -73,7 +72,7 @@ def _apply_observability_id_header(
     response.headers[OBSERVABILITY_ID_HEADER] = observability_id
 
 
-def _matched_route_template(routes: Iterable[BaseRoute], scope: dict) -> str | None:
+def _matched_route_template(routes: Iterable[BaseRoute], scope: Scope) -> str | None:
     for route in routes:
         match, _ = route.matches(scope)
         if match is Match.FULL:
@@ -95,14 +94,14 @@ def _matched_route_template(routes: Iterable[BaseRoute], scope: dict) -> str | N
     return None
 
 
-def _native_route_template(app: FastAPI, scope: dict) -> str | None:
+def _native_route_template(app: FastAPI, scope: Scope) -> str | None:
     """Return the matched FastAPI route template, excluding the Flask mount."""
 
     return _matched_route_template(app.router.routes, scope)
 
 
 def create_asgi_app(
-    wsgi_app: WSGIApp,
+    wsgi_app: Any,  # noqa: ANN401  # Flask's WSGI signature and a2wsgi's WSGIApp stubs disagree
     *,
     route_settings: RouteImplementationSettings | None = None,
     dependencies: NativeRouteDependencies | None = None,
diff --git a/policyengine_api/constants.py b/policyengine_api/constants.py
index 6aeef3fda..cea6ecd7d 100644
--- a/policyengine_api/constants.py
+++ b/policyengine_api/constants.py
@@ -3,6 +3,7 @@
 from datetime import datetime
 from importlib.metadata import distribution, distributions
 from pathlib import Path
+from typing import Any
 
 from policyengine_api.dataset_display import (
     DEFAULT_DATASET_DISPLAY_LABEL,
@@ -47,7 +48,7 @@ def _resolve_distribution_version(
     return "0.0.0"
 
 
-def get_py_manifest() -> Path:
+def get_py_manifest() -> Any:  # noqa: ANN401  # SimplePath, which open() rejects; a Path at runtime
     return distribution("policyengine").locate_file(
         "policyengine/data/bundle/manifest.json"
     )
diff --git a/policyengine_api/country.py b/policyengine_api/country.py
index 93300b936..332871fae 100644
--- a/policyengine_api/country.py
+++ b/policyengine_api/country.py
@@ -9,9 +9,8 @@
 from typing import Any
 
 import dpath
-import numpy as np
 import pandas as pd
-from policyengine_core.model_api import Enum, Reform
+from policyengine_core.model_api import Enum, Reform, Variable
 from policyengine_core.parameters import (
     Parameter,
     ParameterNode,
@@ -23,7 +22,7 @@
 from policyengine_core.populations import Population
 from policyengine_core.simulations import Simulation
 from policyengine_core.taxbenefitsystems import TaxBenefitSystem
-from policyengine_core.variables import Variable
+from policyengine_core.types import ArrayLike
 
 from policyengine_api.constants import (
     get_bundle_default_dataset_option,
@@ -66,7 +65,7 @@ def _serialize_float(value: object) -> float | str:
 
 
 def _serialize_axis_result(
-    result: np.ndarray,
+    result: ArrayLike,
     variable: Variable,
     entity_index: int,
     count_entities: int,
@@ -74,7 +73,9 @@ def _serialize_axis_result(
     if variable.value_type == Enum:
         result = result.decode_to_str()
 
-    entity_values = result.reshape((-1, count_entities)).T[entity_index].tolist()
+    entity_values: list[object] = (
+        result.reshape((-1, count_entities)).T[entity_index].tolist()
+    )
 
     if variable.value_type is float:
         return [_serialize_float(value) for value in entity_values]
@@ -153,7 +154,10 @@ def build_metadata(self) -> None:
             )
         )
 
-    def _json_safe(self, value: object) -> object:
+    def _json_safe(
+        self,
+        value: object,
+    ) -> Any:  # noqa: ANN401  # other values pass through, assumed JSON-safe
         if isinstance(value, Path):
             return str(value)
         if isinstance(value, (date, datetime)):
diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py
index 7fe178e6f..7eefd896e 100644
--- a/policyengine_api/routes/household_routes.py
+++ b/policyengine_api/routes/household_routes.py
@@ -139,7 +139,9 @@ def _spm_error_response(error: Exception) -> Response | None:
 
 
 def _calculation_response(calculation: HouseholdCalculationResult) -> dict:
-    result = dict(status="ok", message=None, result=calculation.household)
+    result: dict[str, object] = dict(
+        status="ok", message=None, result=calculation.household
+    )
     if calculation.warnings:
         result["warnings"] = list(calculation.warnings)
     if calculation.spm_provenance is not None:
diff --git a/policyengine_api/services/economy_comparison.py b/policyengine_api/services/economy_comparison.py
index f0f992ee5..08de03777 100644
--- a/policyengine_api/services/economy_comparison.py
+++ b/policyengine_api/services/economy_comparison.py
@@ -336,9 +336,9 @@ def intra_decile_impact(baseline: dict, reform: dict) -> dict:
     for lower, upper, label in zip(BOUNDS[:-1], BOUNDS[1:], LABELS):
         outcome_groups[label] = []
         for i in range(1, 11):
-            in_decile: bool = decile == i
-            in_group: bool = (income_change > lower) & (income_change <= upper)
-            in_both: bool = in_decile & in_group
+            in_decile: np.ndarray = decile == i
+            in_group: np.ndarray = (income_change > lower) & (income_change <= upper)
+            in_both: np.ndarray = in_decile & in_group
 
             people_in_both: np.float64 = people[in_both].sum()
             people_in_decile: np.float64 = people[in_decile].sum()
@@ -388,9 +388,9 @@ def intra_wealth_decile_impact(baseline: dict, reform: dict) -> dict:
     for lower, upper, label in zip(BOUNDS[:-1], BOUNDS[1:], LABELS):
         outcome_groups[label] = []
         for i in range(1, 11):
-            in_decile: bool = decile == i
-            in_group: bool = (income_change > lower) & (income_change <= upper)
-            in_both: bool = in_decile & in_group
+            in_decile: np.ndarray = decile == i
+            in_group: np.ndarray = (income_change > lower) & (income_change <= upper)
+            in_both: np.ndarray = in_decile & in_group
 
             people_in_both: np.float64 = people[in_both].sum()
             people_in_decile: np.float64 = people[in_decile].sum()
diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py
index 5979682c8..4c72feac6 100644
--- a/policyengine_api/services/economy_service.py
+++ b/policyengine_api/services/economy_service.py
@@ -919,7 +919,7 @@ def _get_or_create_economic_impact(
             severity="INFO",
         )
 
-        most_recent_impact: dict | None = self._get_most_recent_impact(
+        most_recent_impact: CachedReformImpact | None = self._get_most_recent_impact(
             setup_options=setup_options
         )
         if most_recent_impact and self._should_refresh_cached_impact(
@@ -1180,7 +1180,7 @@ def _get_previous_impacts(
     def _get_most_recent_impact(
         self,
         setup_options: EconomicImpactSetupOptions,
-    ) -> dict | None:
+    ) -> CachedReformImpact | None:
         """
         Get the first impact from the reform_impacts table based on the provided setup options.
         Returns the first impact if it exists, otherwise returns None.
@@ -1207,7 +1207,7 @@ def _get_most_recent_impact(
 
     def _determine_impact_action(
         self,
-        most_recent_impact: ReformImpact | None,
+        most_recent_impact: CachedReformImpact | None,
     ) -> ImpactAction:
         if not most_recent_impact:
             return ImpactAction.CREATE
@@ -1552,7 +1552,7 @@ def _build_options_hash_lookup_pattern(self, options_hash: str) -> str:
             return f"{escaped_options_hash[:-1]}&%"
         return f"{escaped_options_hash}%"
 
-    def _extract_cached_result(self, most_recent_impact: ReformImpact) -> dict:
+    def _extract_cached_result(self, most_recent_impact: CachedReformImpact) -> dict:
         try:
             return self._parse_json_object(most_recent_impact.reform_impact_json)
         except (TypeError, ValueError):
@@ -1561,7 +1561,7 @@ def _extract_cached_result(self, most_recent_impact: ReformImpact) -> dict:
     def _should_refresh_cached_impact(
         self,
         setup_options: EconomicImpactSetupOptions,
-        most_recent_impact: ReformImpact,
+        most_recent_impact: CachedReformImpact,
     ) -> bool:
         if most_recent_impact.status == ImpactStatus.COMPUTING.value:
             return False
@@ -1599,7 +1599,7 @@ def _should_refresh_cached_impact(
 
     def _with_policyengine_bundle(
         self,
-        result: dict,
+        result: dict | None,
         setup_options: EconomicImpactSetupOptions,
         execution: ModalSimulationExecution | None = None,
     ) -> dict:
diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py
index 1e67ec09c..2957d5bd9 100644
--- a/policyengine_api/services/household_service.py
+++ b/policyengine_api/services/household_service.py
@@ -35,13 +35,11 @@
 
     from sqlalchemy.orm import Session, sessionmaker
 
-    from policyengine_api.spm import SPMSelection
-
 
 def household_storage_json(
     country_id: str,
     household_json: dict,
-    spm: SPMSelection | dict | None = None,
+    spm: object = None,
 ) -> dict:
     """Keep the measurement selection in the same atomic JSON/hash as inputs.
 
@@ -164,7 +162,7 @@ def create_household(
         label: str | None,
         *,
         record_mirror_event: bool = False,
-        spm: dict | None = None,
+        spm: object = None,
     ) -> HouseholdCreateResult:
         # Validate before the persistence boundary so typed input failures remain 400.
         household_json = household_storage_json(country_id, household_json, spm)
diff --git a/policyengine_api/specification.py b/policyengine_api/specification.py
index adc5d6631..f4ae4f5a2 100644
--- a/policyengine_api/specification.py
+++ b/policyengine_api/specification.py
@@ -4,7 +4,7 @@
 
 import json
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import yaml
 
@@ -35,7 +35,9 @@ def _openapi_30_schema(value: JSONValue) -> JSONValue:
         return [_openapi_30_schema(item) for item in value]
     if not isinstance(value, dict):
         return value
-    result = {key: _openapi_30_schema(item) for key, item in value.items()}
+    result: dict[str, Any] = {
+        key: _openapi_30_schema(item) for key, item in value.items()
+    }
     if "const" in result:
         result["enum"] = [result.pop("const")]
     variants = result.get("anyOf", [])
@@ -72,7 +74,9 @@ def _economy_query_parameters(model: type[BaseModel]) -> list[dict[str, object]]
     return result
 
 
-def _component_schema(model: type[BaseModel]) -> JSONValue:
+def _component_schema(
+    model: Any,  # noqa: ANN401  # mypy joins the model-class dict to ModelMetaclass
+) -> JSONValue:
     """Generate one OpenAPI component and publish nested models separately."""
     schema = model.model_json_schema(
         ref_template="#/components/schemas/{model}",
diff --git a/policyengine_api/spm.py b/policyengine_api/spm.py
index 2dda6d2f5..9919aa729 100644
--- a/policyengine_api/spm.py
+++ b/policyengine_api/spm.py
@@ -476,7 +476,7 @@ def _selected_forecast(expected_sha256: str) -> SPMForecast:
 
 def normalize_spm_selection(
     country_id: str,
-    selection: SPMSelection | dict | None,
+    selection: object,
     *,
     stored: bool = False,
 ) -> dict | None:
diff --git a/policyengine_api/utils/cache_utils.py b/policyengine_api/utils/cache_utils.py
index 87d74b66c..cbd02da9d 100644
--- a/policyengine_api/utils/cache_utils.py
+++ b/policyengine_api/utils/cache_utils.py
@@ -3,6 +3,7 @@
 import hashlib
 import json
 import logging
+from typing import Any
 
 import flask
 
@@ -18,7 +19,7 @@ def make_cache_key(*args: object, **kwargs: object) -> str:
     different across workers/restarts; that made same-input requests
     miss the cache in production.
     """
-    data = ""
+    data: Any = ""
     if flask.request.content_type == "application/json":
         data = flask.request.get_json()
     elif flask.request.content_type in [
diff --git a/policyengine_api/utils/json.py b/policyengine_api/utils/json.py
index 9b617b95a..54f54e8d7 100644
--- a/policyengine_api/utils/json.py
+++ b/policyengine_api/utils/json.py
@@ -2,7 +2,7 @@
 
 import base64
 import hashlib
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -10,7 +10,9 @@
     from policyengine_api.json_types import JSONValue
 
 
-def make_hashable(o: object) -> object:
+def make_hashable(
+    o: object,
+) -> Any:  # noqa: ANN401  # returns scalars unchanged, and callers sort them
     if isinstance(o, (tuple, list)):
         return tuple(make_hashable(e) for e in o)
 

From 4f0bd70cec4dad818a3079cc8ea7abf4bec78224 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:33:48 +0400
Subject: [PATCH 08/10] Annotate the .github helper scripts for ruff ANN

bump_version.py, find-api-model-versions.py and scripts/check_updates.py
(with a ChangelogEntry TypedDict). Annotations only.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .github/bump_version.py            |  4 ++--
 .github/find-api-model-versions.py |  2 +-
 .github/scripts/check_updates.py   | 31 +++++++++++++++++++++---------
 3 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/.github/bump_version.py b/.github/bump_version.py
index 05092fdad..1e7b49e18 100644
--- a/.github/bump_version.py
+++ b/.github/bump_version.py
@@ -47,7 +47,7 @@ def bump_version(version: str, bump: str) -> str:
     return f"{major}.{minor}.{patch + 1}"
 
 
-def update_file(path: Path, old_version: str, new_version: str):
+def update_file(path: Path, old_version: str, new_version: str) -> None:
     text = path.read_text()
     updated = text.replace(
         f'version = "{old_version}"',
@@ -58,7 +58,7 @@ def update_file(path: Path, old_version: str, new_version: str):
         print(f"  Updated {path}")
 
 
-def main():
+def main() -> None:
     root = Path(__file__).resolve().parent.parent
     pyproject = root / "pyproject.toml"
     changelog_dir = root / "changelog.d"
diff --git a/.github/find-api-model-versions.py b/.github/find-api-model-versions.py
index ece882b19..0c0096db2 100644
--- a/.github/find-api-model-versions.py
+++ b/.github/find-api-model-versions.py
@@ -88,7 +88,7 @@ def find_api_model_versions() -> dict[str, str]:
     }
 
 
-def find_api_model_versions_and_output_to_github():
+def find_api_model_versions_and_output_to_github() -> None:
     """
     Find the API model versions and output them to a file for GitHub.
     """
diff --git a/.github/scripts/check_updates.py b/.github/scripts/check_updates.py
index 951f143b2..eccbc982e 100644
--- a/.github/scripts/check_updates.py
+++ b/.github/scripts/check_updates.py
@@ -13,6 +13,7 @@
 import argparse
 import re
 import sys
+from typing import TypedDict
 
 import requests
 
@@ -23,7 +24,12 @@
 }
 
 
-def fetch_changelog(package):
+class ChangelogEntry(TypedDict):
+    version: str
+    changes: dict[str, list[str]]
+
+
+def fetch_changelog(package: str) -> str | None:
     """Fetch CHANGELOG.md from the package's GitHub repo."""
     repo = REPO_MAP.get(package)
     if not repo:
@@ -36,18 +42,18 @@ def fetch_changelog(package):
     return None
 
 
-def parse_version(version_str):
+def parse_version(version_str: str) -> tuple[int, ...]:
     """Parse a version string into a comparable tuple."""
     return tuple(map(int, version_str.split(".")))
 
 
-def parse_changelog(text):
+def parse_changelog(text: str | None) -> list[ChangelogEntry]:
     """Parse Keep-a-Changelog markdown into structured entries."""
     if not text:
         return []
 
-    entries = []
-    current_entry = None
+    entries: list[ChangelogEntry] = []
+    current_entry: ChangelogEntry | None = None
     current_category = None
 
     for line in text.splitlines():
@@ -77,7 +83,9 @@ def parse_changelog(text):
     return entries
 
 
-def get_changes_between(changelog, old_version, new_version):
+def get_changes_between(
+    changelog: list[ChangelogEntry], old_version: str, new_version: str
+) -> list[ChangelogEntry]:
     """Extract changelog entries between two versions."""
     old_v = parse_version(old_version)
     new_v = parse_version(new_version)
@@ -88,9 +96,14 @@ def get_changes_between(changelog, old_version, new_version):
     ]
 
 
-def format_changes(entries):
+def format_changes(entries: list[ChangelogEntry]) -> str:
     """Format changelog entries as markdown sections."""
-    buckets = {"added": [], "changed": [], "fixed": [], "removed": []}
+    buckets: dict[str, list[str]] = {
+        "added": [],
+        "changed": [],
+        "fixed": [],
+        "removed": [],
+    }
     for entry in entries:
         for cat, items in entry.get("changes", {}).items():
             if cat in buckets:
@@ -105,7 +118,7 @@ def format_changes(entries):
     return "\n\n".join(sections) if sections else ""
 
 
-def main():
+def main() -> int:
     parser = argparse.ArgumentParser(
         description="Fetch changelog entries between two package versions."
     )

From ccfa83b10b90a367e8f3fe2fcc5562030a45660c Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:34:08 +0400
Subject: [PATCH 09/10] Register ANN and the test-file ruff exemptions

api.quality.ruff-lint now includes ANN, and the **/tests/** per-file
ignore list (S101, D1, ANN, PLR2004) is registered as
api.quality.tests-ruff-exemptions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 AGENTS.md                               |  7 ++-----
 changelog.d/ruff-annotations.changed.md |  1 +
 docs/generated/rules.md                 | 17 +++++++++++------
 rules/registry.toml                     | 21 ++++++++++++++++++++-
 4 files changed, 34 insertions(+), 12 deletions(-)
 create mode 100644 changelog.d/ruff-annotations.changed.md

diff --git a/AGENTS.md b/AGENTS.md
index 3415e4f33..e2e3342c5 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -32,10 +32,6 @@ names, logging, or observability failure handling, read
 When modifying the `Makefile`, read
 `docs/engineering/skills/repository-maintenance.md`.
 
-When adding a Pydantic, SQLModel or SQLAlchemy model class that other models
-inherit from, or moving an import under `if TYPE_CHECKING:`, read
-`docs/engineering/skills/typing-only-imports.md`.
-
 When adding or changing a GitHub Actions workflow step or a script under
 `.github/scripts/`, read `docs/engineering/skills/ci-workflows.md`.
 
@@ -126,9 +122,10 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`.
 - `api.quality.environment-variables-load`: Required environment variables and credentials load in a deployed-like environment. (blocking: CI job `test_env_vars`)
 - `api.quality.mypy-v2`: Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. (blocking: mypy)
 - `api.quality.ruff-format`: All Python files are formatted with ruff format. (blocking: ruff format)
-- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W))
+- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W))
 - `api.quality.runtime-evaluated-models-resolve`: Every package model class that other models inherit from is listed in ruff's runtime-evaluated-base-classes, and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. (blocking: pytest `tests/unit/test_runtime_evaluated_annotations.py`)
 - `api.quality.tests-pass`: The unit, contract and to-refactor test suites pass. (blocking: CI job `test`)
+- `api.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**`)
 
 ### Observability
 
diff --git a/changelog.d/ruff-annotations.changed.md b/changelog.d/ruff-annotations.changed.md
new file mode 100644
index 000000000..01854b08e
--- /dev/null
+++ b/changelog.d/ruff-annotations.changed.md
@@ -0,0 +1 @@
+Enabled ruff's annotation rules (ANN) and annotated every function outside tests; test files are exempt from ANN, missing docstrings, magic numbers and bare-assert rules.
diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index d155a209b..2674efff8 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -8,11 +8,11 @@ Every rule that governs `policyengine-api` is listed here with how it is enforce
 
 | Tier | enforced | ratcheted | advisory | doc-only | Total |
 | --- | ---: | ---: | ---: | ---: | ---: |
-| mechanical | 42 | 0 | 1 | 12 | 55 |
+| mechanical | 43 | 0 | 1 | 12 | 56 |
 | metric | 0 | 0 | 0 | 0 | 0 |
 | judgment | 0 | 0 | 0 | 0 | 0 |
 
-**Blocking:** 42 of 55 rules.
+**Blocking:** 43 of 56 rules.
 
 ## Gaps
 
@@ -269,9 +269,10 @@ flowchart LR
 | `api.quality.environment-variables-load` | Required environment variables and credentials load in a deployed-like environment. | mechanical | enforced | yes | CI job `test_env_vars` | `pr.yml:test_env_vars` | — |
 | `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — |
 | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — |
-| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — |
+| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — |
 | `api.quality.runtime-evaluated-models-resolve` | Every package model class that other models inherit from is listed in ruff's runtime-evaluated-base-classes, and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** |
 | `api.quality.tests-pass` | The unit, contract and to-refactor test suites pass. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — |
+| `api.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**` | `pr.yml:lint` | — |
 
 ```mermaid
 flowchart LR
@@ -288,12 +289,14 @@ flowchart LR
   e_ruff_format["ruff format"]
   j_CI_pr_yml_lint["CI pr.yml:lint"]
   r_api_quality_ruff_lint("api.quality.ruff-lint (enforced)")
-  e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"]
+  e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"]
   r_api_quality_runtime_evaluated_models_resolve("api.quality.runtime-evaluated-models-resolve (enforced)")
   e_pytest_tests_unit_test_runtime_evaluated_annotations_py["pytest tests/unit/test_runtime_evaluated_annotations.py"]
   j_CI_pr_yml_test["CI pr.yml:test"]
   r_api_quality_tests_pass("api.quality.tests-pass (enforced)")
   e_CI_job_test["CI job test"]
+  r_api_quality_tests_ruff_exemptions("api.quality.tests-ruff-exemptions (enforced)")
+  e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___["ruff per-file ignore of ANN, D1, PLR2004, S101 for **/tests/**"]
   r_api_quality_alembic_revisions_unused_imports --> e_ruff_per_file_ignore_of_ALL_for_migrations___versions___py
   e_ruff_per_file_ignore_of_ALL_for_migrations___versions___py -.-> local_only
   r_api_quality_environment_variables_load --> e_CI_job_test_env_vars
@@ -302,12 +305,14 @@ flowchart LR
   e_mypy --> j_CI_pr_yml_typecheck_v2
   r_api_quality_ruff_format --> e_ruff_format
   e_ruff_format --> j_CI_pr_yml_lint
-  r_api_quality_ruff_lint --> e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_
-  e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint
+  r_api_quality_ruff_lint --> e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_
+  e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint
   r_api_quality_runtime_evaluated_models_resolve --> e_pytest_tests_unit_test_runtime_evaluated_annotations_py
   e_pytest_tests_unit_test_runtime_evaluated_annotations_py --> j_CI_pr_yml_test
   r_api_quality_tests_pass --> e_CI_job_test
   e_CI_job_test --> j_CI_pr_yml_test
+  r_api_quality_tests_ruff_exemptions --> e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___
+  e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___ --> j_CI_pr_yml_lint
 ```
 
 ## Observability
diff --git a/rules/registry.toml b/rules/registry.toml
index 46b8c6490..556bb1ad8 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -679,7 +679,7 @@ autofix = true
 [rule.enforcer]
 kind = "ruff"
 mode = "lint"
-codes = ["E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"]
+codes = ["ANN", "E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"]
 [[rule.ci]]
 workflow = ".github/workflows/pr.yml"
 job = "lint"
@@ -722,6 +722,25 @@ mode = "per-file-ignore"
 pattern = "migrations/*/versions/*.py"
 codes = ["ALL"]
 
+[[rule]]
+id = "api.quality.tests-ruff-exemptions"
+category = "quality"
+statement = "Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests."
+rationale = "docs/engineering/skills/rules-registry.md#baselines"
+scope = ["tests/**/*.py", "scripts/rules/tests/**/*.py"]
+tier = "mechanical"
+status = "enforced"
+autofix = false
+[rule.enforcer]
+kind = "ruff"
+mode = "per-file-ignore"
+pattern = "**/tests/**"
+codes = ["ANN", "D1", "PLR2004", "S101"]
+[[rule.ci]]
+workflow = ".github/workflows/pr.yml"
+job = "lint"
+invocation = "ruff check ."
+
 [[rule]]
 id = "api.quality.mypy-v2"
 category = "quality"

From 4b11c794de0bbbb4a2490837ef95c6072efaace7 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Thu, 15 Oct 2026 19:32:30 +0400
Subject: [PATCH 10/10] Resolve model bases the way ruff does in the
 runtime-annotation guard

Ruff treats a class as runtime-evaluated when a base resolves, through
the import in the class's own module, to a listed name, or when a base
defined in the same module is itself runtime-evaluated. It matches the
import path, not where the base is defined, so a listed base imported
through a re-export is not recognised. The guard compared definition
paths, so it would have passed such a subclass while ruff asked for its
annotation imports to be moved.

The guard now parses each model's module, resolves its bases through
the module's imports (including relative imports, attribute access and
subscripted generics), follows same-module bases, and fails with the
path to add. test_base_resolution_matches_ruff pins the resolution:
listed imports, module attributes, same-module chains, generics and
relative imports count; re-exports and unlisted third-party bases do
not. An experiment against ruff 0.15.12 confirmed each case. No model
in the package is affected today. The rule statement and the
typing-only-imports skill describe the matching precisely.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 AGENTS.md                                     |   2 +-
 changelog.d/ruff-annotations.changed.md       |   2 +-
 .../engineering/skills/typing-only-imports.md |  34 ++-
 docs/generated/rules.md                       |   2 +-
 rules/registry.toml                           |   3 +-
 .../test_runtime_evaluated_annotations.py     | 201 +++++++++++++++---
 6 files changed, 202 insertions(+), 42 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index e2e3342c5..543b529fc 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -123,7 +123,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`.
 - `api.quality.mypy-v2`: Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. (blocking: mypy)
 - `api.quality.ruff-format`: All Python files are formatted with ruff format. (blocking: ruff format)
 - `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W))
-- `api.quality.runtime-evaluated-models-resolve`: Every package model class that other models inherit from is listed in ruff's runtime-evaluated-base-classes, and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. (blocking: pytest `tests/unit/test_runtime_evaluated_annotations.py`)
+- `api.quality.runtime-evaluated-models-resolve`: Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. (blocking: pytest `tests/unit/test_runtime_evaluated_annotations.py`)
 - `api.quality.tests-pass`: The unit, contract and to-refactor test suites pass. (blocking: CI job `test`)
 - `api.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**`)
 
diff --git a/changelog.d/ruff-annotations.changed.md b/changelog.d/ruff-annotations.changed.md
index 01854b08e..1aceb6e5b 100644
--- a/changelog.d/ruff-annotations.changed.md
+++ b/changelog.d/ruff-annotations.changed.md
@@ -1 +1 @@
-Enabled ruff's annotation rules (ANN) and annotated every function outside tests; test files are exempt from ANN, missing docstrings, magic numbers and bare-assert rules.
+Enabled ruff's annotation rules (ANN) and annotated every function outside tests; test files are exempt from ANN, missing docstrings, magic numbers and bare-assert rules. The runtime-annotation guard now resolves each model's bases through its module's imports, as ruff does, so a base imported through a re-export is caught.
diff --git a/docs/engineering/skills/typing-only-imports.md b/docs/engineering/skills/typing-only-imports.md
index f3e66b4e8..b442abf5b 100644
--- a/docs/engineering/skills/typing-only-imports.md
+++ b/docs/engineering/skills/typing-only-imports.md
@@ -21,17 +21,28 @@ Some libraries do evaluate those strings at runtime:
 - FastAPI reads endpoint signatures when routes are registered.
 
 Names those annotations use must exist at runtime. Ruff keeps a class's
-annotation imports at runtime only when one of the class's **direct** bases is
-listed in `runtime-evaluated-base-classes`, or when an endpoint's decorator is
-listed in `runtime-evaluated-decorators`, under `[tool.ruff.lint.flake8-type-checking]`
-in `pyproject.toml`. Ruff does not follow inheritance, so every package class
-that other models inherit from is listed by name.
+annotation imports at runtime only when it treats the class as runtime-evaluated,
+or when an endpoint's decorator is listed in `runtime-evaluated-decorators`, under
+`[tool.ruff.lint.flake8-type-checking]` in `pyproject.toml`. A class is
+runtime-evaluated when one of its bases:
+
+- resolves, through the import in the class's own module, to a name listed in
+  `runtime-evaluated-base-classes`; or
+- is a class defined in the same module that is itself runtime-evaluated.
+
+Ruff matches the path the module imports the base from, not where the base is
+defined. A listed base imported through a re-export (`from package import Base`,
+where `package/__init__.py` re-exports it) does not match its listed definition
+path. Ruff does not follow inheritance across modules, so every package class
+that models in other modules inherit from is listed.
 
 ## Adding a model base
 
-When you add a model class that other models inherit from, add its fully
-qualified name (module where it is defined, then class name) to
-`runtime-evaluated-base-classes` in the same change.
+When you add a model class that models in other modules inherit from, add the
+qualified name those modules import it by (normally the module where it is
+defined, then the class name) to `runtime-evaluated-base-classes` in the same
+change. Import model bases from the module that defines them, so one listed path
+covers every subclass.
 
 If ruff reports a typing-only import in a model, do not move the import until you
 have checked that the model's base is listed. Moving it for a model whose base
@@ -39,6 +50,7 @@ is missing breaks the model at runtime: Pydantic raises "not fully defined" the
 first time the model is used, and SQLAlchemy fails to configure its mappers.
 
 `tests/unit/test_runtime_evaluated_annotations.py` enforces both halves
-(`api.quality.runtime-evaluated-models-resolve`): every package model base is
-listed, and every Pydantic model and SQLAlchemy mapper resolves its annotations
-once the package is imported.
+(`api.quality.runtime-evaluated-models-resolve`): every package model is one ruff
+treats as runtime-evaluated, resolving its bases the way ruff does, and every
+Pydantic model and SQLAlchemy mapper resolves its annotations once the package is
+imported.
diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 2674efff8..4778b566a 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -270,7 +270,7 @@ flowchart LR
 | `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — |
 | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — |
 | `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — |
-| `api.quality.runtime-evaluated-models-resolve` | Every package model class that other models inherit from is listed in ruff's runtime-evaluated-base-classes, and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** |
+| `api.quality.runtime-evaluated-models-resolve` | Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** |
 | `api.quality.tests-pass` | The unit, contract and to-refactor test suites pass. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — |
 | `api.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**` | `pr.yml:lint` | — |
 
diff --git a/rules/registry.toml b/rules/registry.toml
index 556bb1ad8..08096ff8e 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -688,7 +688,7 @@ invocation = "ruff check ."
 [[rule]]
 id = "api.quality.runtime-evaluated-models-resolve"
 category = "quality"
-statement = "Every package model class that other models inherit from is listed in ruff's runtime-evaluated-base-classes, and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported."
+statement = "Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported."
 rationale = "docs/engineering/skills/typing-only-imports.md#models-that-read-annotations-at-runtime"
 scope = ["policyengine_api/**/*.py", "pyproject.toml"]
 tier = "mechanical"
@@ -698,6 +698,7 @@ autofix = false
 kind = "pytest"
 nodes = [
   "tests/unit/test_runtime_evaluated_annotations.py::test_every_model_base_is_listed_for_ruff",
+  "tests/unit/test_runtime_evaluated_annotations.py::test_base_resolution_matches_ruff",
   "tests/unit/test_runtime_evaluated_annotations.py::test_every_pydantic_model_resolves_its_annotations",
   "tests/unit/test_runtime_evaluated_annotations.py::test_every_sqlalchemy_mapper_configures",
   "tests/unit/test_runtime_evaluated_annotations.py::test_skipped_app_modules_define_no_classes",
diff --git a/tests/unit/test_runtime_evaluated_annotations.py b/tests/unit/test_runtime_evaluated_annotations.py
index 2c75d57ea..5f9fb6bcf 100644
--- a/tests/unit/test_runtime_evaluated_annotations.py
+++ b/tests/unit/test_runtime_evaluated_annotations.py
@@ -3,9 +3,12 @@
 The package's modules use ``from __future__ import annotations``, so ruff's
 typing-only import rules (TC) move imports used only in annotations under
 ``if TYPE_CHECKING:``. Pydantic, SQLModel and SQLAlchemy read model annotations
-at runtime, so ruff keeps a class's annotation imports at runtime only when one
-of its direct bases is listed in ``runtime-evaluated-base-classes`` in
-pyproject.toml. See docs/engineering/skills/typing-only-imports.md.
+at runtime, so ruff keeps a class's annotation imports at runtime only when it
+treats the class as runtime-evaluated: one of its bases resolves, through the
+import in the class's own module, to a name listed in
+``runtime-evaluated-base-classes`` in pyproject.toml, or is a class in the same
+module that is itself runtime-evaluated. Ruff does not follow re-exports. See
+docs/engineering/skills/typing-only-imports.md.
 """
 
 from __future__ import annotations
@@ -13,7 +16,9 @@
 import ast
 import importlib
 import pkgutil
+import sys
 import tomllib
+from dataclasses import dataclass
 from functools import cache
 from pathlib import Path
 
@@ -32,6 +37,117 @@
 APP_ASSEMBLY_MODULES = frozenset({"policyengine_api.api", "policyengine_api.asgi"})
 MODEL_ROOTS = (pydantic.BaseModel, sqlmodel.SQLModel, DeclarativeBase)
 
+# A module that exercises each way a base can resolve, for
+# test_base_resolution_matches_ruff.
+EXAMPLE_MODULE = """\
+import pydantic
+from typing import Generic, TypeVar
+from policyengine_api.models import Listed
+from policyengine_api.reexports import Listed as ReexportedListed
+from pydantic_settings import BaseSettings
+from .siblings import Sibling
+class ViaImport(Listed): ...
+class ViaModuleAttribute(pydantic.BaseModel): ...
+class LocalBase(Listed): ...
+class ViaSameModuleChain(LocalBase): ...
+class ViaGeneric(Listed, Generic[TypeVar("T")]): ...
+class ViaReexport(ReexportedListed): ...
+class ViaUnlistedThirdParty(BaseSettings): ...
+class ViaRelativeImport(Sibling): ...
+"""
+
+
+@dataclass(frozen=True)
+class ModuleSource:
+    """A module's top-level classes and the qualified names its imports bind."""
+
+    name: str
+    classes: dict[str, ast.ClassDef]
+    imported_names: dict[str, str]
+
+
+def _parse_module_source(name: str, text: str, *, is_package: bool) -> ModuleSource:
+    tree = ast.parse(text)
+    package = name if is_package else name.rpartition(".")[0]
+    imported_names: dict[str, str] = {}
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ImportFrom):
+            if node.level:
+                anchor = (
+                    package.rsplit(".", node.level - 1)[0]
+                    if node.level > 1
+                    else package
+                )
+                module = f"{anchor}.{node.module}" if node.module else anchor
+            else:
+                module = node.module or ""
+            for alias in node.names:
+                imported_names[alias.asname or alias.name] = f"{module}.{alias.name}"
+        elif isinstance(node, ast.Import):
+            for alias in node.names:
+                if alias.asname:
+                    imported_names[alias.asname] = alias.name
+                else:
+                    top_level = alias.name.split(".")[0]
+                    imported_names[top_level] = top_level
+    classes = {node.name: node for node in tree.body if isinstance(node, ast.ClassDef)}
+    return ModuleSource(name=name, classes=classes, imported_names=imported_names)
+
+
+@cache
+def _module_source(module_name: str) -> ModuleSource:
+    path = Path(sys.modules[module_name].__file__)
+    return _parse_module_source(
+        module_name,
+        path.read_text(encoding="utf-8"),
+        is_package=path.name == "__init__.py",
+    )
+
+
+def _qualified_base_name(source: ModuleSource, base: ast.expr) -> str | None:
+    """Resolve a base expression the way ruff does: through the module's imports."""
+    if isinstance(base, ast.Subscript):
+        base = base.value
+    attributes = []
+    while isinstance(base, ast.Attribute):
+        attributes.append(base.attr)
+        base = base.value
+    if not isinstance(base, ast.Name):
+        return None
+    if base.id in source.imported_names:
+        head = source.imported_names[base.id]
+    elif base.id in source.classes:
+        head = f"{source.name}.{base.id}"
+    else:
+        head = base.id
+    return ".".join([head, *reversed(attributes)])
+
+
+def _is_runtime_evaluated(
+    source: ModuleSource,
+    class_node: ast.ClassDef,
+    listed: frozenset[str],
+    visiting: frozenset[str] = frozenset(),
+) -> bool:
+    """Return whether ruff treats the class's annotations as read at runtime."""
+    for base in class_node.bases:
+        base_name = _qualified_base_name(source, base)
+        if base_name in listed:
+            return True
+        local_name = (base_name or "").removeprefix(f"{source.name}.")
+        if (
+            local_name in source.classes
+            and local_name not in visiting
+            and _is_runtime_evaluated(
+                source,
+                source.classes[local_name],
+                listed,
+                visiting | {class_node.name},
+            )
+        ):
+            return True
+    return False
+
 
 @cache
 def _import_package_modules() -> tuple[str, ...]:
@@ -51,8 +167,7 @@ def _generic_origin(model: type) -> type:
     """Return the class a Pydantic generic like ``Model[Item]`` was built from.
 
     Pydantic creates ``Model[Item]`` at runtime as a subclass of ``Model``. Ruff
-    never sees it as a class of its own: a ``class Child(Model[Item])`` statement
-    names ``Model``.
+    never sees it as a class of its own.
     """
     metadata = getattr(model, "__pydantic_generic_metadata__", None)
     origin = metadata.get("origin") if metadata else None
@@ -79,11 +194,24 @@ def _package_model_classes() -> list[type]:
     )
 
 
-def _listed_runtime_evaluated_bases() -> set[str]:
+@cache
+def _listed_runtime_evaluated_bases() -> frozenset[str]:
     with (PROJECT_ROOT / "pyproject.toml").open("rb") as pyproject:
         settings = tomllib.load(pyproject)
     type_checking = settings["tool"]["ruff"]["lint"]["flake8-type-checking"]
-    return set(type_checking["runtime-evaluated-base-classes"])
+    return frozenset(type_checking["runtime-evaluated-base-classes"])
+
+
+def _not_runtime_evaluated(model: type) -> str | None:
+    """Describe a top-level model ruff would not treat as runtime-evaluated."""
+    source = _module_source(model.__module__)
+    class_node = source.classes.get(model.__qualname__)
+    if class_node is None or _is_runtime_evaluated(
+        source, class_node, _listed_runtime_evaluated_bases()
+    ):
+        return None
+    bases = [_qualified_base_name(source, base) for base in class_node.bases]
+    return f"{_qualified_name(model)}: bases as ruff resolves them: {bases}"
 
 
 def _unresolved_annotation(model: type[pydantic.BaseModel]) -> str | None:
@@ -95,28 +223,47 @@ def _unresolved_annotation(model: type[pydantic.BaseModel]) -> str | None:
 
 
 def test_every_model_base_is_listed_for_ruff() -> None:
-    models = _package_model_classes()
-    package_models = set(models)
-    inherited_bases = {
-        _generic_origin(base)
-        for model in models
-        for base in model.__bases__
-        if _generic_origin(base) in package_models
-    }
-    listed = _listed_runtime_evaluated_bases()
-    missing = sorted(
-        _qualified_name(base)
-        for base in inherited_bases
-        if _qualified_name(base) not in listed
-    )
-    assert missing == [], (
-        "these package models are inherited by other models but are missing from "
+    unlisted = [
+        message
+        for model in _package_model_classes()
+        if (message := _not_runtime_evaluated(model)) is not None
+    ]
+    assert unlisted == [], (
+        "ruff would not treat these models' annotations as read at runtime, so it "
+        "would ask for imports they need to be moved under TYPE_CHECKING. Add the "
+        "base path ruff resolves (the path the module imports it from) to "
         "[tool.ruff.lint.flake8-type-checking] runtime-evaluated-base-classes in "
-        "pyproject.toml, so ruff would move imports their subclasses need at "
-        "runtime; add them:\n" + "\n".join(missing)
+        "pyproject.toml:\n" + "\n".join(unlisted)
     )
 
 
+def test_base_resolution_matches_ruff() -> None:
+    source = _parse_module_source(
+        "policyengine_api.example", EXAMPLE_MODULE, is_package=False
+    )
+    listed = frozenset(
+        {
+            "policyengine_api.models.Listed",
+            "pydantic.BaseModel",
+            "policyengine_api.siblings.Sibling",
+        }
+    )
+    runtime_evaluated = {
+        name
+        for name, class_node in source.classes.items()
+        if _is_runtime_evaluated(source, class_node, listed)
+    }
+
+    assert runtime_evaluated == {
+        "ViaImport",
+        "ViaModuleAttribute",
+        "LocalBase",
+        "ViaSameModuleChain",
+        "ViaGeneric",
+        "ViaRelativeImport",
+    }
+
+
 def test_every_pydantic_model_resolves_its_annotations() -> None:
     unresolved = [
         message
@@ -126,8 +273,8 @@ def test_every_pydantic_model_resolves_its_annotations() -> None:
     ]
     assert unresolved == [], (
         "these models use names in annotations that don't exist at runtime, "
-        "usually an import moved under TYPE_CHECKING for a model whose base is "
-        "missing from runtime-evaluated-base-classes:\n" + "\n".join(unresolved)
+        "usually an import moved under TYPE_CHECKING for a model ruff doesn't treat "
+        "as runtime-evaluated:\n" + "\n".join(unresolved)
     )