Enable ruff annotation rules (ANN)

.github/scripts/update_policyengine_package.py

@@ -28,7 +28,7 @@ def _mapping(value: object) -> Mapping[str, Any]:
     return value if isinstance(value, Mapping) else {}
 
 
-def _json_from_stdin() -> Any:
+def _json_from_stdin() -> object:
     return json.load(sys.stdin)
 
 

AGENTS.md

@@ -143,9 +143,10 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`.
 - `sim.quality.every-project-tested`: Every project and library with unit tests runs them in the PR test matrix. (doc-only, not checked: follow it yourself)
 - `sim.quality.ruff-format-src`: Source code under projects/*/src and libs/*/src is formatted with ruff format. (blocking: ruff format)
 - `sim.quality.ruff-format-tests`: Test code is formatted with ruff format too; CI currently checks only src directories. (doc-only, not checked: follow it yourself)
-- `sim.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in ruff.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))
+- `sim.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in ruff.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))
 - `sim.quality.runtime-evaluated-models-resolve`: Every Pydantic model is runtime-evaluated for ruff: one of its bases, read through the imports of the model's own module, is listed in ruff.toml's runtime-evaluated-base-classes or is a runtime-evaluated class in that same module; and every Pydantic model resolves its annotations once its package is imported. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py`, `projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py`)
 - `sim.quality.tests-pass`: Each project's unit tests pass in its own locked environment. (blocking: CI job `test`)
+- `sim.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/**`)
 
 ### CI workflows and scripts
 

changelog_entry.yaml

@@ -4,6 +4,7 @@
       - Made CLAUDE.md and the Copilot instructions thin pointers to AGENTS.md, required workflow logic longer than one line to live in tested and shellchecked scripts, and registered the tracking-issue rule
       - Added a rules registry (rules/registry.toml) with a standard-library tool that validates it, checks every rule is enforced where it claims, and renders docs/generated/rules.md and the Rules block in AGENTS.md
     changed:
+      - 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 typing-only import rules (TC) and moved annotation-only imports under TYPE_CHECKING, keeping runtime imports for Pydantic models and FastAPI endpoints
       - Listed policyengine.py's Simulation, pydantic-settings' BaseSettings and the contract's SingleYearMacroOutput as runtime-evaluated bases for ruff, imported SingleYearMacroOutput from the contract in the executor's compat models instead of through a re-export, and added executor and entrypoint tests that fail when ruff, reading a model's bases through its module's imports, would not treat a Pydantic model as runtime-evaluated, or when a model's annotations do not resolve
       - Ran ruff check in CI with the first group of rule families (imports, pyupgrade and other safe auto-fixes) and fixed every finding

docs/generated/rules.md

@@ -8,11 +8,11 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 
 | Tier | enforced | ratcheted | advisory | doc-only | Total |
 | --- | ---: | ---: | ---: | ---: | ---: |
-| mechanical | 34 | 1 | 0 | 14 | 49 |
+| mechanical | 35 | 1 | 0 | 14 | 50 |
 | metric | 0 | 0 | 0 | 0 | 0 |
 | judgment | 0 | 0 | 0 | 0 | 0 |
 
-**Blocking:** 35 of 49 rules.
+**Blocking:** 36 of 50 rules.
 
 ## Gaps
 
@@ -249,9 +249,10 @@ flowchart LR
 | `sim.quality.every-project-tested` | Every project and library with unit tests runs them in the PR test matrix. | mechanical | doc-only | no | nothing | — | — |
 | `sim.quality.ruff-format-src` | Source code under projects/*/src and libs/*/src is formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — |
 | `sim.quality.ruff-format-tests` | Test code is formatted with ruff format too; CI currently checks only src directories. | mechanical | doc-only | no | nothing | — | — |
-| `sim.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in ruff.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` | — |
+| `sim.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in ruff.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` | — |
 | `sim.quality.runtime-evaluated-models-resolve` | Every Pydantic model is runtime-evaluated for ruff: one of its bases, read through the imports of the model's own module, is listed in ruff.toml's runtime-evaluated-base-classes or is a runtime-evaluated class in that same module; and every Pydantic model resolves its annotations once its package is imported. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py`, `projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py` | `pr.yml:test` | yes |
 | `sim.quality.tests-pass` | Each project's unit tests pass in its own locked environment. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — |
+| `sim.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
@@ -262,22 +263,26 @@ flowchart LR
   j_CI_pr_yml_lint["CI pr.yml:lint"]
   r_sim_quality_ruff_format_tests("sim.quality.ruff-format-tests (doc-only)")
   r_sim_quality_ruff_lint("sim.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_sim_quality_runtime_evaluated_models_resolve("sim.quality.runtime-evaluated-models-resolve (enforced)")
   e_pytest_projects_policyengine_simulation_entry_tests_test_runtime_evaluated_annotations_py__projects_policyengine_simulation_executor_tests_test_runtime_evaluated_annotations_py["pytest projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py, projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py"]
   j_CI_pr_yml_test["CI pr.yml:test"]
   r_sim_quality_tests_pass("sim.quality.tests-pass (enforced)")
   e_CI_job_test["CI job test"]
+  r_sim_quality_tests_ruff_exemptions("sim.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_sim_quality_every_project_tested -.-> gap_unenforced
   r_sim_quality_ruff_format_src --> e_ruff_format
   e_ruff_format --> j_CI_pr_yml_lint
   r_sim_quality_ruff_format_tests -.-> gap_unenforced
-  r_sim_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_sim_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_sim_quality_runtime_evaluated_models_resolve --> e_pytest_projects_policyengine_simulation_entry_tests_test_runtime_evaluated_annotations_py__projects_policyengine_simulation_executor_tests_test_runtime_evaluated_annotations_py
   e_pytest_projects_policyengine_simulation_entry_tests_test_runtime_evaluated_annotations_py__projects_policyengine_simulation_executor_tests_test_runtime_evaluated_annotations_py --> j_CI_pr_yml_test
   r_sim_quality_tests_pass --> e_CI_job_test
   e_CI_job_test --> j_CI_pr_yml_test
+  r_sim_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
 ```
 
 ## CI workflows and scripts

libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py

@@ -16,7 +16,7 @@ class JWTDecoder:
     the parsed token with all fields.
     """
 
-    def __init__(self, issuer: str, audience: str, auto_error: bool = True):
+    def __init__(self, issuer: str, audience: str, auto_error: bool = True) -> None:
         self.issuer = issuer
         self.audience = audience
         self.auto_error = auto_error
@@ -38,7 +38,7 @@ def __call__(
             LOG.info("Ignoring invalid auth token")
         return None
 
-    def _validate(self, token: HTTPAuthorizationCredentials | None):
+    def _validate(self, token: HTTPAuthorizationCredentials | None) -> dict[str, str]:
         if token is None:
             LOG.info("No bearer token on request")
             raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)

libs/policyengine-fastapi/src/policyengine_fastapi/database.py

@@ -30,7 +30,7 @@ def create_session_dep(engine: Engine) -> SessionGeneratorFactory:
     routers.
     """
 
-    def session_dep():
+    def session_dep() -> Generator[Session]:
         with Session(engine) as session:
             yield session
 

libs/policyengine-fastapi/src/policyengine_fastapi/exit.py

@@ -4,7 +4,7 @@
 """
 
 import logging
-from collections.abc import Callable
+from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from typing import Any, TypeVar
 
@@ -14,24 +14,24 @@
 
 
 class AppExit:
-    def __init__(self):
+    def __init__(self) -> None:
         self.callbacks: list[Callable[[], None]] = []
 
-    def __call__(self, *args: Any, **kwargs: Any):
+    def __call__(self, *args: object, **kwargs: object) -> Callable[[T], T]:
         def accept(func: T) -> T:
             self.callbacks.append(lambda: func(*args, **kwargs))
             return func
 
         return accept
 
-    def _exit(self):
+    def _exit(self) -> None:
         log.info("Executing exit callbacks")
         for cb in self.callbacks:
             cb()
         log.info("Exit callbacks executed")
 
     @asynccontextmanager
-    async def lifespan(self):
+    async def lifespan(self) -> AsyncIterator[None]:
         yield
         self._exit()
 

libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py

@@ -24,7 +24,7 @@ class HealthStatus(BaseModel):
 
 
 class HealthSystemReporter:
-    def __init__(self, name: str, probes: dict[str, HealthProbe]):
+    def __init__(self, name: str, probes: dict[str, HealthProbe]) -> None:
         self.name = name
         self.probes = probes
 
@@ -45,10 +45,10 @@ class HealthRegistry:
 
     systems: dict[str, HealthSystemReporter]
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.systems = {}
 
-    def register(self, reporter: HealthSystemReporter):
+    def register(self, reporter: HealthSystemReporter) -> None:
         self.systems[reporter.name] = reporter
 
     def report(self) -> HealthStatus:

libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from typing import Self
 
 from .stages import TracerCaptureMode
 
@@ -44,5 +45,5 @@ class ObservabilityConfig:
     slow_run_threshold_seconds: float = 30.0
 
     @classmethod
-    def disabled(cls, service_name: str = "policyengine-observability"):
+    def disabled(cls, service_name: str = "policyengine-observability") -> Self:
         return cls(enabled=False, service_name=service_name)

libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py

@@ -2,15 +2,14 @@
 
 import json
 from hashlib import sha256
-from typing import Any
 from uuid import uuid4
 
 
 def generate_observability_id() -> str:
     return str(uuid4())
 
 
-def stable_config_hash(payload: Any) -> str:
+def stable_config_hash(payload: object) -> str:
     serialised = json.dumps(
         payload,
         sort_keys=True,

libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py

@@ -2,24 +2,30 @@
 
 from contextlib import AbstractContextManager
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol, Self
 
 from .config import ObservabilityConfig
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
+    from types import TracebackType
 
     from .contracts import SimulationLifecycleEvent, TracerArtifactManifest
 
 
 class NoOpSpan(AbstractContextManager["NoOpSpan"]):
-    def __enter__(self):
+    def __enter__(self) -> Self:
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
         return None
 
-    def set_attribute(self, key: str, value: Any) -> None:
+    def set_attribute(self, key: str, value: object) -> None:
         return None
 
     def add_event(self, name: str, attributes: Mapping[str, Any] | None = None) -> None:

libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py

@@ -19,7 +19,7 @@
 # Configure opentelemetry
 # 1. to include python logs and
 # 2. export to console (for demo purposes)
-def export_ot_to_console(resource: Resource):
+def export_ot_to_console(resource: Resource) -> None:
     """
     configure opentelemetry to dump messages to console for debugging on desktop
     """
@@ -33,6 +33,6 @@ def export_ot_to_console(resource: Resource):
     metrics.set_meter_provider(meterProvider)
 
     @exit()
-    def flush_ot():
+    def flush_ot() -> None:
         traceProvider.force_flush()
         meterProvider.force_flush()

libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py

@@ -23,11 +23,11 @@
 
 
 class AddGcpProjectToTraceFilter(logging.Filter):
-    def __init__(self, project: str):
+    def __init__(self, project: str) -> None:
         super().__init__()
         self.project = project
 
-    def filter(self, record: logging.LogRecord):
+    def filter(self, record: logging.LogRecord) -> bool | logging.LogRecord:
         if not hasattr(record, "otelTraceID"):
             return super().filter(record)
         record.otelTraceID = f"projects/{self.project}/traces/{record.otelTraceID}"  # type: ignore
@@ -58,15 +58,15 @@ def _get_project_id_from_metadata() -> str | None:
 
 
 class GCPLoggingInstrumentor:
-    def __init__(self, project_id: str | None = None):
+    def __init__(self, project_id: str | None = None) -> None:
         self.project_id = project_id or _get_project_id()
 
     """
     Configures the standard opentelemetry logging instrumentor to generate
     json logs as per gcp expectations.
     """
 
-    def instrument(self):
+    def instrument(self) -> None:
         from opentelemetry.instrumentation.logging import LoggingInstrumentor
         from pythonjsonlogger.json import JsonFormatter
 
@@ -93,7 +93,7 @@ def instrument(self):
         LoggingInstrumentor().instrument()
 
 
-def export_ot_to_gcp(resource: Resource):
+def export_ot_to_gcp(resource: Resource) -> None:
     from opentelemetry.exporter.cloud_monitoring import (
         CloudMonitoringMetricsExporter,
     )
@@ -114,7 +114,7 @@ def export_ot_to_gcp(resource: Resource):
     metrics.set_meter_provider(meterProvider)
 
     @exit()
-    def flush_ot():
+    def flush_ot() -> None:
         log.info("Flushing traces and metrics before shutdown")
         traceProvider.force_flush()
         meterProvider.force_flush()

libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py

@@ -10,7 +10,7 @@ class FastAPIEnhancedInstrumenter:
     metrics instead of global for the whole api.
     """
 
-    def instrument(self, app: FastAPI):
+    def instrument(self, app: FastAPI) -> None:
         FastAPIInstrumentor.instrument_app(app)
         middleware = Middleware(app)
         app.middleware("http")(middleware)

libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py

@@ -1,10 +1,10 @@
 from dataclasses import dataclass
 from time import time
-from typing import Any
 
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Response
 from fastapi.routing import APIRoute
 from opentelemetry import metrics
+from starlette.middleware.base import RequestResponseEndpoint
 from starlette.routing import Match
 
 
@@ -16,7 +16,7 @@ class Route:
 
 
 class Middleware:
-    def __init__(self, app: FastAPI):
+    def __init__(self, app: FastAPI) -> None:
         self.app = app
         self.meter = metrics.get_meter("policyengine")
 
@@ -32,7 +32,9 @@ def _create_route(self, route: APIRoute) -> Route:
             latency=self.meter.create_histogram(f"{name}_latency"),
         )
 
-    async def __call__(self, request: Request, call_next) -> Any:
+    async def __call__(
+        self, request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
         route: Route | None = None
         for r in self.routes:
             if r.route.matches(request.scope)[0] != Match.NONE:

libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py

@@ -24,14 +24,14 @@ async def ping(request: PingRequest) -> PingResponse:
         return PingResponse(incremented=request.value + 1)
 
     @router.get("/ping/started", response_model=str)
-    async def started():
+    async def started() -> str:
         """
         Verify the api is running.
         """
         return "alive"
 
     @router.get("/ping/alive", response_model=HealthStatus)
-    async def alive():
+    async def alive() -> JSONResponse:
         """
         Check if the service is healthy. This will always return
         data, but will return a 503 response code if the service is
@@ -46,5 +46,5 @@ async def alive():
     return router
 
 
-def include_all_routers(api: FastAPI, health_registry: HealthRegistry):
+def include_all_routers(api: FastAPI, health_registry: HealthRegistry) -> None:
     api.include_router(create_router(health_registry))

libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py

@@ -25,11 +25,11 @@
 BUDGET_WINDOW_JOB_SEED_DICT_NAME = "simulation-api-budget-window-job-seeds"
 
 
-def _budget_window_job_store():
+def _budget_window_job_store() -> modal.Dict:
     return modal.Dict.from_name(BUDGET_WINDOW_JOB_DICT_NAME, create_if_missing=True)
 
 
-def _budget_window_job_seed_store():
+def _budget_window_job_seed_store() -> modal.Dict:
     return modal.Dict.from_name(
         BUDGET_WINDOW_JOB_SEED_DICT_NAME, create_if_missing=True
     )