diff --git a/.github/scripts/update_policyengine_package.py b/.github/scripts/update_policyengine_package.py index 1d3ab3ce6..56197deb3 100644 --- a/.github/scripts/update_policyengine_package.py +++ b/.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) diff --git a/AGENTS.md b/AGENTS.md index d5cb893ea..0bae03c1c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -138,8 +138,9 @@ 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.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 diff --git a/changelog_entry.yaml b/changelog_entry.yaml index fd5edd93f..e4d918c70 100644 --- a/changelog_entry.yaml +++ b/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 - Ran ruff check in CI with the first group of rule families (imports, pyupgrade and other safe auto-fixes) and fixed every finding - Moved the multi-line `run` blocks of the simulation deployment workflows into single-line steps and tested scripts without changing their behaviour or outputs, leaving only the release-sentinel and routing-rollback blocks in the one-line `run` rule's baseline diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 6dac8864c..1d979899d 100644 --- a/docs/generated/rules.md +++ b/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 | 32 | 1 | 0 | 14 | 47 | +| mechanical | 33 | 1 | 0 | 14 | 48 | | metric | 0 | 0 | 0 | 0 | 0 | | judgment | 0 | 0 | 0 | 0 | 0 | -**Blocking:** 33 of 47 rules. +**Blocking:** 34 of 48 rules. ## Gaps @@ -245,8 +245,9 @@ 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.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 @@ -257,18 +258,22 @@ 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_tests_pass("sim.quality.tests-pass (enforced)") e_CI_job_test["CI job test"] j_CI_pr_yml_test["CI pr.yml: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_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 diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py index 07331be90..a8ec9b929 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py +++ b/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) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/database.py b/libs/policyengine-fastapi/src/policyengine_fastapi/database.py index e5569f372..d1340a26e 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/database.py +++ b/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 diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py b/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py index e2821e764..5f7cde3c6 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py +++ b/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() diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py index 464826daf..66c9c64bd 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py +++ b/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: diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py index fbe58ebd3..465cf198a 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py +++ b/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) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py index 1dc46157e..4282d0f3b 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py @@ -2,7 +2,6 @@ import json from hashlib import sha256 -from typing import Any from uuid import uuid4 @@ -10,7 +9,7 @@ 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, diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py index 3a06f307d..05d4d7d3d 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py +++ b/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: diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py index dc265d56e..31b21b294 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py +++ b/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() diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py index 968fed712..91f40ed8f 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py +++ b/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,7 +58,7 @@ 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() """ @@ -66,7 +66,7 @@ def __init__(self, project_id: str | None = None): 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() diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py index c91265317..aab1ce0aa 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py +++ b/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) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py index e39ea23e0..f741b51f9 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py +++ b/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: diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py index bd53cb5ac..ce927959e 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py +++ b/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)) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py index 85d3cddf4..1e396ca0b 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py +++ b/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 ) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index f4580ed20..83e9a6ba8 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py @@ -41,7 +41,7 @@ class PolicyParameterChanges(RootModel[JsonObject]): """Policy parameter names mapped to their dated values.""" -def _move_internal_telemetry_alias(value): +def _move_internal_telemetry_alias(value: object) -> object: if not isinstance(value, dict): return value if "_telemetry" in value and "telemetry" not in value: @@ -50,7 +50,7 @@ def _move_internal_telemetry_alias(value): return value -def _strip_internal_passthrough_fields(value): +def _strip_internal_passthrough_fields(value: object) -> object: """Drop internal-only fields the gateway adds to payloads in flight. Parent batch entrypoints attach ``_metadata`` describing resolved routing, @@ -71,7 +71,7 @@ def _strip_internal_passthrough_fields(value): } -def _enforce_max_payload_size(value): +def _enforce_max_payload_size(value: object) -> object: if not isinstance(value, dict): return value try: @@ -126,17 +126,17 @@ class GatewayRequestBase(BaseModel): @model_validator(mode="before") @classmethod - def move_internal_telemetry_alias(cls, value): + def move_internal_telemetry_alias(cls, value: object) -> object: return _move_internal_telemetry_alias(value) @model_validator(mode="before") @classmethod - def strip_internal_passthrough_fields(cls, value): + def strip_internal_passthrough_fields(cls, value: object) -> object: return _strip_internal_passthrough_fields(value) @model_validator(mode="before") @classmethod - def enforce_max_payload_size(cls, value): + def enforce_max_payload_size(cls, value: object) -> object: return _enforce_max_payload_size(value) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index 135bed7b8..160175fa6 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -1,9 +1,9 @@ """Dependency-light public models for US SPM selection and calculation receipts.""" import re -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import date -from typing import Any, Literal +from typing import Any, Literal, Self from pydantic import ( BaseModel, @@ -17,8 +17,10 @@ model_validator, ) +from policyengine_simulation_contract.json_types import JsonObject -def _selection_schema(schema): + +def _selection_schema(schema: dict[str, Any]) -> None: # An omitted option inherits the selected bundle's default. Advertising # Python attribute defaults here makes generated clients send an explicit # county selection even when their caller omitted geography entirely. @@ -45,7 +47,7 @@ class SPMSelection(BaseModel): as_of: str | None = None @model_serializer(mode="wrap") - def serialize_selection( + def serialize_selection( # noqa: ANN201 # a return annotation would become the serialization JSON schema self, handler: SerializerFunctionWrapHandler, info: SerializationInfo ): """Preserve explicitly selected options through ordinary and nested JSON. @@ -75,14 +77,14 @@ def serialize_selection( @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: if 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: # Request options inherit bundle defaults. Validate coupled options # once both are explicit (including after defaults are resolved). if not {"geography_kind", "geography_id"} <= self.model_fields_set: @@ -108,7 +110,7 @@ class SPMResolvedConfiguration(BaseModel): as_of: str | None @model_serializer(mode="wrap") - def serialize_resolved_configuration( + def serialize_resolved_configuration( # noqa: ANN201 # a return annotation would become the serialization JSON schema self, handler: SerializerFunctionWrapHandler, info: SerializationInfo ): """Keep required nullable fields in completed response bodies.""" @@ -158,7 +160,7 @@ class SPMProvenance(BaseModel): runtime_versions: "SPMRuntimeVersions" @model_serializer(mode="wrap") - def serialize_provenance( + def serialize_provenance( # noqa: ANN201 # a return annotation would become the serialization JSON schema self, handler: SerializerFunctionWrapHandler, info: SerializationInfo ): """Keep required nullable configuration fields in response bodies.""" @@ -325,15 +327,15 @@ class SPMErrorDetail(BaseModel): class SPMInputError(ValueError): """Transportable error shared by the control plane and worker.""" - def __init__(self, code: str, message: str): + def __init__(self, code: str, message: str) -> None: self.code = code self.message = message super().__init__(code, message) - def __str__(self): + def __str__(self) -> str: return self.message - def to_dict(self): + def to_dict(self) -> dict[str, str]: return {"code": self.code, "message": self.message} @@ -356,7 +358,7 @@ class SPMCapability(BaseModel): defaults: SPMSelection @model_validator(mode="after") - def pinned_defaults(self): + def pinned_defaults(self) -> Self: if not self.defaults.forecast_content_sha256 or not self.defaults.scenario: raise ValueError( "Certified SPM defaults must pin artifact hash and scenario" @@ -423,14 +425,14 @@ def build_spm_comparison_provenance( def resolve_spm_selection( - country, - selection, + country: str, + selection: SPMSelection | Mapping[str, object] | None, *, - capability, - policyengine_version, - model_version, - route_provenance=None, -): + capability: SPMCapability | Mapping[str, object] | None, + policyengine_version: str | None, + model_version: str | None, + route_provenance: str | None = None, +) -> JsonObject | None: """Resolve only certified metadata; unrecognized future bundles fail closed.""" if country.lower() != "us": if selection is not None: @@ -500,8 +502,8 @@ def resolve_spm_selection( def validate_spm_result( - result: dict, selection: Any, *, expected_year: int | str | None = None -): + result: dict, selection: object, *, expected_year: int | str | None = None +) -> SPMComparisonProvenance | None: """Do not accept incomplete or mixed-method child/cached output.""" if "spm_config" in result: raise SPMInputError( diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index a048b8f3a..98178064d 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -186,7 +186,7 @@ def publish_v2_manifest( class V2ManifestLoader: """Load v2 state while retaining only this loader's last valid v2 copy.""" - def __init__(self, store: V2ManifestStore): + def __init__(self, store: V2ManifestStore) -> None: self._store = store self._last_valid: V2VersionManifest | None = None diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py index dfe54d716..e60ded480 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_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,13 +13,13 @@ def generate_observability_id() -> str: return str(uuid4()) -def resolve_observability_id(value: Any) -> str: +def resolve_observability_id(value: object) -> str: """Use a valid caller value or create a new diagnostic identifier.""" return normalize_observability_id(value) or generate_observability_id() -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. Observability metadata is best effort. Callers must never reject or alter a diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py index 5e1886752..26015dc0a 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py @@ -60,7 +60,7 @@ class TelemetryEnvelope(BaseModel): @model_validator(mode="before") @classmethod - def discard_noncanonical_context_fields(cls, value: Any) -> Any: + def discard_noncanonical_context_fields(cls, value: object) -> object: """Ignore old body fields without using them as correlation inputs. This keeps the receiver compatible with the deployed API during the diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py index fb6ca6e32..031a81330 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py @@ -45,6 +45,8 @@ from datetime import datetime from uuid import UUID + from sqlalchemy import ColumnElement, RowMapping, Table + REPORT_IDENTITY_COLUMNS = ( "environment", "calculation_flow", @@ -155,7 +157,7 @@ def open_connection_once( return connect() -def _database_value(value: Any) -> Any: +def _database_value(value: object) -> object: if isinstance(value, StrEnum): return value.value if isinstance(value, tuple): @@ -168,15 +170,17 @@ def _record_values(record: object) -> dict[str, Any]: return {name: _database_value(value) for name, value in dumped.items()} -def _report(row: Any) -> ComparisonReportRecord: +def _report(row: RowMapping) -> ComparisonReportRecord: return ComparisonReportRecord.model_validate(dict(row)) -def _simulation(row: Any) -> ComparisonSimulationRecord: +def _simulation(row: RowMapping) -> ComparisonSimulationRecord: return ComparisonSimulationRecord.model_validate(dict(row)) -def _identity_predicates(table: Any, values: dict[str, Any], columns: tuple[str, ...]): +def _identity_predicates( + table: Table, values: dict[str, Any], columns: tuple[str, ...] +) -> tuple[ColumnElement[bool], ...]: return tuple(table.c[name] == values[name] for name in columns) diff --git a/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py b/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py index 183b50d9f..1d9609e0d 100644 --- a/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py +++ b/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py @@ -47,7 +47,7 @@ def access_token() -> str: return token -def test_protected_route_rejects_missing_token(base_url: str): +def test_protected_route_rejects_missing_token(base_url: str) -> None: response = httpx.post( f"{base_url}{PROTECTED_PATH}", json=PROBE_PAYLOAD, @@ -57,7 +57,7 @@ def test_protected_route_rejects_missing_token(base_url: str): assert response.status_code == 403 -def test_protected_route_rejects_invalid_token(base_url: str): +def test_protected_route_rejects_invalid_token(base_url: str) -> None: response = httpx.post( f"{base_url}{PROTECTED_PATH}", json=PROBE_PAYLOAD, @@ -71,7 +71,7 @@ def test_protected_route_rejects_invalid_token(base_url: str): def test_valid_token_reaches_authenticated_old_gateway( base_url: str, access_token: str, -): +) -> None: response = httpx.post( f"{base_url}{PROTECTED_PATH}", json=PROBE_PAYLOAD, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index ba77098c9..506905399 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -8,7 +8,7 @@ import uuid from collections.abc import Callable from contextlib import asynccontextmanager -from typing import Any, Literal, Protocol +from typing import TYPE_CHECKING, Any, Literal, Protocol from uuid import UUID from fastapi import Depends, FastAPI, Request @@ -68,6 +68,11 @@ TemporaryStage12UnsupportedRequest, ) +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from starlette.middleware.base import RequestResponseEndpoint + logger = logging.getLogger(__name__) BACKEND_RESPONSE_HEADER = { "X-PolicyEngine-Simulation-Backend": "old_gateway", @@ -154,7 +159,7 @@ def create_app( runtime_settings = settings or Settings.from_env() @asynccontextmanager - async def lifespan(_: FastAPI): + async def lifespan(_: FastAPI) -> AsyncIterator[None]: runtime_settings.validate() await runtime_backend.start() try: @@ -239,7 +244,9 @@ def adopt_backend_observability_id( ) @app.middleware("http") - async def request_context(request: Request, call_next): + async def request_context( + request: Request, call_next: RequestResponseEndpoint + ) -> Response: headers = MutableHeaders(scope=request.scope) request_id = headers.get(REQUEST_ID_HEADER) or str(uuid.uuid4()) incoming_observability_id = normalize_observability_id( diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py index 334084a1d..ca6f143f6 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py @@ -24,7 +24,7 @@ class JWTDecoder: """Validate Auth0 RS256 access tokens without database-layer dependencies.""" - def __init__(self, issuer: str, audience: str): + def __init__(self, issuer: str, audience: str) -> None: self.issuer = issuer self.audience = audience self.jwks_client = jwt.PyJWKClient(f"{issuer}.well-known/jwks.json") @@ -75,7 +75,7 @@ def _decoder(issuer: str, audience: str) -> JWTDecoder: class CallerAuthenticator: """FastAPI dependency that preserves the gateway's JWT contract.""" - def __init__(self, settings: Settings, runtime: ObservabilityRuntime): + def __init__(self, settings: Settings, runtime: ObservabilityRuntime) -> None: self.settings = settings self.runtime = runtime diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py index 6cbdcbefd..a74f2ef70 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py @@ -77,7 +77,7 @@ class ClientCredentialsTokenProvider: REFRESH_MARGIN_SECONDS = 60 - def __init__(self, settings: Settings, client: httpx.AsyncClient): + def __init__(self, settings: Settings, client: httpx.AsyncClient) -> None: self.settings = settings self.client = client self._token: str | None = None @@ -132,7 +132,7 @@ def __init__( *, transport: httpx.AsyncBaseTransport | None = None, runtime: ObservabilityRuntime | None = None, - ): + ) -> None: self.settings = settings self.transport = transport self.observability = runtime diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index 8095728cc..771a97148 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -22,7 +22,10 @@ if TYPE_CHECKING: from policyengine_simulation_contract.stage12_bundle import CountryId - from policyengine_simulation_contract.stage12_manifest import V2WorkerVersion + from policyengine_simulation_contract.stage12_manifest import ( + V2CountryWorker, + V2WorkerVersion, + ) class ComparisonSkipReason(StrEnum): @@ -75,7 +78,7 @@ def _skip(reason: ComparisonSkipReason) -> AdaptedComparisonRun: return AdaptedComparisonRun(report=None, skip_reason=reason) -def _country_worker(worker: V2WorkerVersion, country: str): +def _country_worker(worker: V2WorkerVersion, country: str) -> V2CountryWorker | None: return next( (item for item in worker.countries if item.country == country), None, diff --git a/projects/policyengine-simulation-executor/fixtures/fake_modal.py b/projects/policyengine-simulation-executor/fixtures/fake_modal.py index 926967191..99bcc636b 100644 --- a/projects/policyengine-simulation-executor/fixtures/fake_modal.py +++ b/projects/policyengine-simulation-executor/fixtures/fake_modal.py @@ -5,83 +5,100 @@ a per-file copy. The fakes record calls; they run nothing. """ +from __future__ import annotations + import sys from types import ModuleType +from typing import TYPE_CHECKING, Self + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + from pathlib import Path + + import pytest class FakeImage: - def __init__(self): - self.calls = [] + def __init__(self) -> None: + self.calls: list[tuple[object, ...]] = [] @classmethod - def debian_slim(cls, python_version): + def debian_slim(cls, python_version: str) -> Self: image = cls() image.calls.append(("debian_slim", python_version)) return image - def pip_install(self, *packages): + def pip_install(self, *packages: str) -> Self: self.calls.append(("pip_install", packages)) return self - def pip_install_from_requirements(self, requirements_txt, **kwargs): + def pip_install_from_requirements( + self, requirements_txt: str, **kwargs: object + ) -> Self: self.calls.append(("pip_install_from_requirements", requirements_txt, kwargs)) return self - def uv_sync(self, uv_project_dir="./", **kwargs): + def uv_sync(self, uv_project_dir: str = "./", **kwargs: object) -> Self: self.calls.append(("uv_sync", uv_project_dir, kwargs)) return self - def run_commands(self, *commands, **kwargs): + def run_commands(self, *commands: str, **kwargs: object) -> Self: self.calls.append(("run_commands", commands, kwargs)) return self - def env(self, env): + def env(self, env: Mapping[str, str]) -> Self: self.calls.append(("env", env)) return self - def add_local_python_source(self, *args, **kwargs): + def add_local_python_source(self, *args: str, **kwargs: object) -> Self: self.calls.append(("add_local_python_source", args, kwargs)) return self - def add_local_file(self, local_path, remote_path, **kwargs): + def add_local_file( + self, local_path: str | Path, remote_path: str, **kwargs: object + ) -> Self: self.calls.append(("add_local_file", local_path, remote_path, kwargs)) return self - def add_local_dir(self, local_path, remote_path, **kwargs): + def add_local_dir( + self, local_path: str | Path, remote_path: str, **kwargs: object + ) -> Self: self.calls.append(("add_local_dir", local_path, remote_path, kwargs)) return self - def run_function(self, function, **kwargs): + def run_function(self, function: Callable[..., object], **kwargs: object) -> Self: self.calls.append(("run_function", function.__name__, kwargs)) return self class FakeSecret: @staticmethod - def from_name(*args, **kwargs): + def from_name(*args: object, **kwargs: object) -> dict[str, object]: return {"args": args, "kwargs": kwargs} class FakeApp: - def __init__(self, name): + def __init__(self, name: str) -> None: self.name = name - self.function_calls = [] + self.function_calls: list[tuple[str, dict[str, object]]] = [] - def function(self, **kwargs): - def decorator(function): + def function[F: Callable[..., object]](self, **kwargs: object) -> Callable[[F], F]: + def decorator(function: F) -> F: self.function_calls.append((function.__name__, kwargs)) return function return decorator - def local_entrypoint(self, **kwargs): - def decorator(function): + def local_entrypoint[F: Callable[..., object]]( + self, **kwargs: object + ) -> Callable[[F], F]: + def decorator(function: F) -> F: return function return decorator -def install_fake_modal(monkeypatch): +def install_fake_modal(monkeypatch: pytest.MonkeyPatch) -> None: modal = ModuleType("modal") modal.Image = FakeImage modal.Secret = FakeSecret diff --git a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py index 0a6c1ca5a..21204db4c 100644 --- a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py +++ b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py @@ -15,10 +15,16 @@ synthetic inputs instead of whichever PolicyEngine release is installed. """ +from __future__ import annotations + from types import SimpleNamespace +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import pytest -def install_identity_stubs(monkeypatch): +def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: """Stub the identity seams with the canonical synthetic version-set. Returns mutable bundle, receipt, and SPM-capability state, plus a handle diff --git a/projects/policyengine-simulation-executor/fixtures/spm_doubles.py b/projects/policyengine-simulation-executor/fixtures/spm_doubles.py index 790e2ac9b..39ac74819 100644 --- a/projects/policyengine-simulation-executor/fixtures/spm_doubles.py +++ b/projects/policyengine-simulation-executor/fixtures/spm_doubles.py @@ -1,9 +1,15 @@ """Small simulation doubles for executor tests that are not model tests.""" +from __future__ import annotations + from types import SimpleNamespace +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Mapping -def installed_spm_selection(): +def installed_spm_selection() -> dict[str, object] | None: """Return the installed US bundle's default SPM selection, if any.""" from policyengine_simulation_executor.spm import runtime_spm_capability @@ -11,7 +17,9 @@ def installed_spm_selection(): return None if capability is None else capability.defaults.model_dump() -def spm_receipt(selection, year="2026"): +def spm_receipt( + selection: Mapping[str, object], year: int | str = "2026" +) -> dict[str, object]: """Build the smallest calculation record accepted by ``SPMProvenance``.""" return { "forecast_id": "test-only", @@ -31,7 +39,11 @@ def spm_receipt(selection, year="2026"): } -def spm_capable_simulation(selection, year="2026", **attributes): +def spm_capable_simulation( + selection: Mapping[str, object] | None, + year: int | str = "2026", + **attributes: object, +) -> SimpleNamespace: """Return a stand-in for a simulation that measured ``selection``.""" if selection is None: return SimpleNamespace(**attributes) diff --git a/projects/policyengine-simulation-executor/fixtures/test_modal_scripts.py b/projects/policyengine-simulation-executor/fixtures/test_modal_scripts.py index 9e54267f9..c25ae0fd2 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/fixtures/test_modal_scripts.py @@ -1,17 +1,23 @@ """Fixtures and shared constants for Modal script tests.""" +from __future__ import annotations + import os import tempfile from pathlib import Path +from typing import TYPE_CHECKING import pytest +if TYPE_CHECKING: + from collections.abc import Iterator + REPO_ROOT = Path(__file__).parent.parent.parent.parent SCRIPTS_DIR = REPO_ROOT / ".github" / "scripts" @pytest.fixture -def temp_github_output(): +def temp_github_output() -> Iterator[str]: """Create a temporary file to simulate ``GITHUB_OUTPUT``.""" with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as file: yield file.name @@ -19,7 +25,7 @@ def temp_github_output(): @pytest.fixture -def temp_github_step_summary(): +def temp_github_step_summary() -> Iterator[str]: """Create a temporary file to simulate ``GITHUB_STEP_SUMMARY``.""" with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md") as file: yield file.name @@ -27,7 +33,7 @@ def temp_github_step_summary(): @pytest.fixture -def all_modal_scripts(): +def all_modal_scripts() -> list[Path]: """Return Modal and full-stack simulation deployment scripts.""" return [ *SCRIPTS_DIR.glob("modal-*.sh"), diff --git a/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py index b593a0be4..ccc591b04 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py @@ -8,20 +8,20 @@ class FakeCollection: - def __init__(self, records): + def __init__(self, records: list[dict[str, object]]) -> None: self.dataframe = pd.DataFrame(records) class FakeModelOutput: - def __init__(self, payload): + def __init__(self, payload: dict[str, object]) -> None: self.payload = payload - def model_dump(self, *, mode): + def model_dump(self, *, mode: str) -> dict[str, object]: assert mode == "json" return self.payload -def fake_analysis(): +def fake_analysis() -> SimpleNamespace: return SimpleNamespace( program_statistics=FakeCollection( [ diff --git a/projects/policyengine-simulation-executor/scripts/dump_package_version.py b/projects/policyengine-simulation-executor/scripts/dump_package_version.py index afa2d0403..de7a603de 100644 --- a/projects/policyengine-simulation-executor/scripts/dump_package_version.py +++ b/projects/policyengine-simulation-executor/scripts/dump_package_version.py @@ -3,7 +3,9 @@ import tomllib -def get_package_version_from_lock(package_name, lock_file_path="poetry.lock"): +def get_package_version_from_lock( + package_name: str, lock_file_path: str = "poetry.lock" +) -> str | None: try: with open(lock_file_path, "rb") as f: lock_data = tomllib.load(f) diff --git a/projects/policyengine-simulation-executor/src/modal/_image_setup.py b/projects/policyengine-simulation-executor/src/modal/_image_setup.py index f8f7183c3..54788bb6a 100644 --- a/projects/policyengine-simulation-executor/src/modal/_image_setup.py +++ b/projects/policyengine-simulation-executor/src/modal/_image_setup.py @@ -7,8 +7,23 @@ so policyengine_simulation_executor is not importable there at all. """ - -def fetch_artifacts(bucket: str, manifest, *, client=None): +# Annotation-only imports sit behind a local TYPE_CHECKING flag (not +# typing.TYPE_CHECKING) because this module must have no module-level +# imports; the annotations that use them are therefore quoted. +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Mapping + from typing import Any + + from google.cloud import storage + + +def fetch_artifacts( + bucket: str, + manifest: "Mapping[str, Any] | None", + *, + client: "storage.Client | None" = None, +) -> None: """Download every manifest-listed artifact into the image data folder. Runs as an image-build layer between ``.env(VERSION_ENV)`` and @@ -155,7 +170,7 @@ def fetch_artifacts(bucket: str, manifest, *, client=None): raise RuntimeError(f"Artifact fetch did not produce expected files: {absent}") -def snapshot_models(): +def snapshot_models() -> None: """Pre-load models at image build time for fast cold starts.""" import logging diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index 4e5786b01..b30670187 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -9,6 +9,7 @@ import os from pathlib import Path +from typing import TYPE_CHECKING from policyengine_simulation_executor.hf_access_validation import ( HF_MODAL_SECRET_NAME, @@ -38,6 +39,9 @@ import modal +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + def _version_from_env_or_local_bundle_package(env_var: str, package: str) -> str: value = os.environ.get(env_var) @@ -241,7 +245,7 @@ def build_base_simulation_image() -> modal.Image: def _configure_modal_observability( *, service_role: str, -): +) -> "ObservabilityRuntime": return init_process_observability( service_name=APP_NAME, platform="modal", @@ -251,7 +255,7 @@ def _configure_modal_observability( ) -def _set_modal_call_attributes(runtime) -> None: +def _set_modal_call_attributes(runtime: "ObservabilityRuntime") -> None: try: runtime.set_context(function_call_id=modal.current_function_call_id()) except Exception: diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py index 764e38674..02464c0dc 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py @@ -37,7 +37,7 @@ def extract_annual_impact( *, simulation_year: str, child_result: dict[str, Any], - spm: Any = None, + spm: object = None, ) -> BudgetWindowAnnualImpact: receipt = validate_spm_result(child_result, spm, expected_year=simulation_year) budget = child_result.get("budget", {}) diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py index 2fe139224..eb0596a9f 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py @@ -34,7 +34,12 @@ import modal if TYPE_CHECKING: + from types import ModuleType + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.gateway_models import ( + BudgetWindowBatchState, + ) # Polling tuning. The runner busy-loops across child FunctionCall.get(timeout=0) # probes; when no child resolved we sleep before the next probe to stop the @@ -54,7 +59,7 @@ def serialize_batch_status( - state, + state: BudgetWindowBatchState, *, runtime: ObservabilityRuntime, ) -> dict[str, Any]: @@ -68,7 +73,7 @@ def load_or_create_batch_state( context: BudgetWindowBatchContext, *, runtime: ObservabilityRuntime, -): +) -> BudgetWindowBatchState: with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD)): state = get_batch_job_seed(context.batch_job_id) if state is None: @@ -92,11 +97,11 @@ def __init__( context: BudgetWindowBatchContext, *, runtime: ObservabilityRuntime, - modal_module=None, + modal_module: ModuleType | None = None, poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, poll_interval_backoff_factor: float = POLL_INTERVAL_BACKOFF_FACTOR, - ): + ) -> None: self.context = context self.runtime = runtime self.modal = modal if modal_module is None else modal_module diff --git a/projects/policyengine-simulation-executor/src/modal/logging_redaction.py b/projects/policyengine-simulation-executor/src/modal/logging_redaction.py index 0325af918..4f2542b51 100644 --- a/projects/policyengine-simulation-executor/src/modal/logging_redaction.py +++ b/projects/policyengine-simulation-executor/src/modal/logging_redaction.py @@ -15,7 +15,7 @@ SENSITIVE_PARAM_KEYS = ("data", "reform", "baseline") -def redact_params_for_logging(params) -> dict: +def redact_params_for_logging(params: object) -> dict: """Return a shallow copy of ``params`` safe for observability attributes. We preserve routing-relevant fields (country, scope, version, time diff --git a/projects/policyengine-simulation-executor/src/modal/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py index 34d3b0833..1314b0f4c 100644 --- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py @@ -121,7 +121,7 @@ def publish_manifest(bucket: str, plan: dict) -> str: @app.local_entrypoint() -def main(force: bool = False): +def main(force: bool = False) -> None: from policyengine_simulation_executor.artifact_store import resolve_bucket_name from policyengine_simulation_executor.precompute import run_precompute diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 6700dedff..a2e25e823 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -46,6 +46,8 @@ import modal if TYPE_CHECKING: + from types import ModuleType + from policyengine_observability import ObservabilityRuntime logger = logging.getLogger(__name__) @@ -144,10 +146,10 @@ def __init__( *, app_name: str, runtime: ObservabilityRuntime, - modal_module=None, + modal_module: ModuleType | None = None, poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, - ): + ) -> None: self.params = params self.app_name = app_name self.runtime = runtime @@ -177,14 +179,16 @@ def run(self) -> dict[str, Any]: child_results = self._collect(handles) return self._reduce(child_results, country_module=country_module) - def _country_module(self): + def _country_module(self) -> ModuleType: from policyengine_simulation_executor.simulation_runtime import ( _country_module, ) return _country_module(self.country) - def _extend_last_group_with_uncovered_states(self, country_module) -> None: + def _extend_last_group_with_uncovered_states( + self, country_module: ModuleType + ) -> None: """Registry state regions missing from the static partition ride in the last group, loudly, so a new state-level region added to the model cannot silently drop its households from national results. @@ -307,7 +311,9 @@ def _cancel_all(self, handles: list[tuple[list[str], Any]]) -> None: except Exception: pass - def _reduce(self, child_results: list[dict], *, country_module) -> dict[str, Any]: + def _reduce( + self, child_results: list[dict], *, country_module: ModuleType + ) -> dict[str, Any]: from policyengine_simulation_executor.simulation_runtime import ( _parse_year, ) diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index 4318b5210..78ecc52a1 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -161,7 +161,7 @@ def smoke_import_executor() -> dict: @app.local_entrypoint() -def main(): +def main() -> None: report = smoke_import_executor.remote() print(report) print("executor image smoke OK") diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py index 4d091f342..b6cf9069b 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py @@ -439,7 +439,7 @@ def validate_routing_plan(plan: object) -> dict: return plan -def _routing_store(environment: str, *, create_if_missing: bool): +def _routing_store(environment: str, *, create_if_missing: bool) -> "modal.Dict | None": try: return modal.Dict.from_name( ROUTING_STATE_DICT_NAME, @@ -501,7 +501,7 @@ def _require_serialized_routing_workflow() -> None: ) -def _unfinished_publication(store) -> dict | None: +def _unfinished_publication(store: "modal.Dict | None") -> dict | None: record = store.get(ROUTING_PUBLICATION_KEY) if store is not None else None if record is None: return None @@ -517,7 +517,9 @@ def _unfinished_publication(store) -> dict | None: return record["plan"] if record["status"] == "pending" else None -def _require_finished_publication(store, *, allowed_plans: Iterable[dict] = ()) -> None: +def _require_finished_publication( + store: "modal.Dict | None", *, allowed_plans: Iterable[dict] = () +) -> None: pending = _unfinished_publication(store) if pending is not None and not any(pending == plan for plan in allowed_plans): raise RoutingConflict( @@ -528,7 +530,7 @@ def _require_finished_publication(store, *, allowed_plans: Iterable[dict] = ()) ) -def _write_publication_record(store, *, plan: dict, status: str) -> None: +def _write_publication_record(store: "modal.Dict", *, plan: dict, status: str) -> None: record = {"status": status, "plan": _json_copy(plan)} store[ROUTING_PUBLICATION_KEY] = record if store.get(ROUTING_PUBLICATION_KEY) != record: @@ -1028,7 +1030,7 @@ def restore_routing_state( return next_state -def main(): +def main() -> None: parser = argparse.ArgumentParser( formatter_class=argparse.RawDescriptionHelpFormatter, description=( @@ -1175,7 +1177,7 @@ def main(): ) args = parser.parse_args() - def active_state(): + def active_state() -> object: snapshot = _modal_dict_snapshot( name=ROUTING_STATE_DICT_NAME, environment=args.environment ) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index f79c96578..dab372b87 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -41,7 +41,10 @@ import modal if TYPE_CHECKING: - from policyengine_simulation_contract.stage12_bundle import CountryId + from policyengine_simulation_contract.stage12_bundle import ( + CountryId, + Stage12CountryBundle, + ) STAGE12_DATA_DIR = "/opt/policyengine/stage12-data" _UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "." @@ -87,7 +90,7 @@ def _external_assertions(environment: dict[str, str]) -> dict[str, str]: ] -def _country_bundle(country: CountryId): +def _country_bundle(country: CountryId) -> Stage12CountryBundle: return next( item for item in RESOLVED_BUNDLE.bundle.countries if item.country == country ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py index 98f1a49b6..433343dff 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py @@ -9,7 +9,7 @@ """ -def initialize(app: FastAPI, runtime: ObservabilityRuntime): +def initialize(app: FastAPI, runtime: ObservabilityRuntime) -> None: """ attach all routes to the app and configure them to use the provided SQLModel engine and jwt settings. diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py index b16b155f8..9298fd237 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py @@ -24,12 +24,15 @@ import os from collections.abc import Mapping from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from google.api_core.exceptions import NotFound, PreconditionFailed from policyengine_simulation_executor import artifact_keys +if TYPE_CHECKING: + from google.cloud import storage + ARTIFACT_BUCKET_ENV = "POLICYENGINE_ARTIFACT_BUCKET" @@ -42,7 +45,7 @@ def resolve_bucket_name(explicit: str | None = None) -> str: return bucket -def _default_client(): +def _default_client() -> storage.Client: from google.cloud import storage from policyengine_simulation_executor.simulation_runtime import ( @@ -56,17 +59,19 @@ def _default_client(): class ArtifactStore: - def __init__(self, bucket_name: str | None = None, *, client: Any = None): + def __init__( + self, bucket_name: str | None = None, *, client: storage.Client | None = None + ) -> None: self.bucket_name = resolve_bucket_name(bucket_name) self._client = client @property - def client(self) -> Any: + def client(self) -> storage.Client: if self._client is None: self._client = _default_client() return self._client - def _blob(self, path: str): + def _blob(self, path: str) -> storage.Blob: return self.client.bucket(self.bucket_name).blob(path) def exists(self, path: str) -> bool: @@ -123,7 +128,9 @@ def download_file(self, path: str, local_path: str | Path) -> None: raise partial.replace(destination) - def _upload_json(self, path: str, payload: Mapping[str, Any], *, overwrite: bool): + def _upload_json( + self, path: str, payload: Mapping[str, Any], *, overwrite: bool + ) -> bool: data = json.dumps(payload, sort_keys=True, indent=2) blob = self._blob(path) if overwrite: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py index e66a59092..ffc8d7091 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py @@ -21,13 +21,16 @@ from __future__ import annotations import logging -from typing import Any +from typing import TYPE_CHECKING, Any from policyengine.core import Simulation from pydantic import PrivateAttr from policyengine_simulation_executor import artifact_keys +if TYPE_CHECKING: + from policyengine.core.scoping_strategy import ScopingStrategy + logger = logging.getLogger(__name__) # Artifact outcomes reported via the `baseline_artifact` observability @@ -45,9 +48,9 @@ def qualifying_baseline_identity( *, country: str, dataset_is_default: bool, - policy: Any, + policy: object, region_code: str | None, - scoping_strategy: Any, + scoping_strategy: ScopingStrategy | None, year: int, ) -> artifact_keys.BaselineArtifactIdentity | None: """Full artifact identity for a qualifying baseline, else None. @@ -103,9 +106,9 @@ def deterministic_baseline_id( *, country: str, dataset_is_default: bool, - policy: Any, + policy: object, region_code: str | None, - scoping_strategy: Any, + scoping_strategy: ScopingStrategy | None, year: int, ) -> str | None: """The runtime wrapper: id for a qualifying baseline, else None. @@ -150,7 +153,7 @@ def artifact_outcome(self) -> str | None: """Set by ``ensure()``; None until then.""" return self._artifact_outcome - def run(self): + def run(self) -> None: self._computed_this_process = True super().run() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py index b0b0d7383..9e8391583 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py @@ -1,4 +1,5 @@ import logging +from collections.abc import AsyncIterator from contextlib import asynccontextmanager from fastapi import FastAPI @@ -25,7 +26,7 @@ @asynccontextmanager -async def lifespan(app: FastAPI): +async def lifespan(app: FastAPI) -> AsyncIterator[None]: async with exit.lifespan(): try: yield diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index f8e2caa78..42f3c47be 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -34,7 +34,7 @@ import logging import time -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from policyengine_simulation_executor.artifact_keys import ( BaselineArtifactIdentity, @@ -56,6 +56,13 @@ if TYPE_CHECKING: from collections.abc import Callable + from pathlib import Path + from types import ModuleType + + from policyengine.core import Simulation + + from policyengine_simulation_executor.artifact_store import ArtifactStore + from policyengine_simulation_executor.simulation_runtime import RegionResolution logger = logging.getLogger(__name__) @@ -85,7 +92,9 @@ def cohort_params(year: int, group: list[str]) -> dict: } -def _cohort_resolution(year: int, group: list[str]): +def _cohort_resolution( + year: int, group: list[str] +) -> tuple[dict[str, Any], ModuleType, RegionResolution]: """Cohort request params, country module, and region resolution. The planner's identity (``cohort_identity``) and the worker's @@ -277,7 +286,9 @@ def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildR ) -def _prepare_cohort_baseline(bucket: str, expected: BaselinePlanEntry): +def _prepare_cohort_baseline( + bucket: str, expected: BaselinePlanEntry +) -> tuple[ArtifactStore, Path, Simulation]: """Shared by compute and verify: dataset download + simulation build.""" from pathlib import Path diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py index a6abc66e2..47dc14793 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py @@ -162,6 +162,6 @@ class DeterminismVerdict(_StrictModel): class RemoteFunction(Protocol): """The slice of a Modal function handle the orchestration uses.""" - def remote(self, *args: Any) -> Any: ... + def remote(self, *args: Any) -> Any: ... # noqa: ANN401 # mirrors modal.Function.remote; payloads and results vary per function - def spawn(self, *args: Any) -> Any: ... + def spawn(self, *args: Any) -> Any: ... # noqa: ANN401 # mirrors modal.Function.spawn; payloads vary and the returned FunctionCall is untyped here diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py index 0d8ce99b5..522b6df3e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py @@ -14,6 +14,7 @@ from dataclasses import dataclass from functools import lru_cache from pathlib import Path +from typing import TYPE_CHECKING from policyengine_simulation_contract.dataset_uri import ( runtime_dataset_uri, @@ -21,6 +22,12 @@ split_dataset_revision, ) +if TYPE_CHECKING: + from policyengine.provenance.manifest import ( + CountryReleaseManifest, + DataPackageVersion, + ) + os.environ.setdefault("POLICYENGINE_SKIP_COUNTRY_IMPORTS", "1") SUPPORTED_COUNTRIES = frozenset({"us", "uk"}) @@ -63,7 +70,7 @@ def _default_dataset_sha256( *, default_dataset: str, data_release: Mapping, - manifest, + manifest: CountryReleaseManifest, ) -> str | None: """The certified SHA-256 of the default dataset, from the bundle. @@ -99,7 +106,7 @@ def _normalise_country(country: str) -> str: return country -def _artifact_revision(data_package) -> str: +def _artifact_revision(data_package: DataPackageVersion) -> str: return data_package.release_manifest_revision or data_package.version @@ -129,7 +136,7 @@ def _with_hf_revision_unvalidated(dataset_uri: str, revision: str) -> str: def _artifact_hf_uri( *, - data_package, + data_package: DataPackageVersion, reference: object, ) -> tuple[str | None, str | None]: from policyengine.provenance.manifest import build_hf_uri @@ -152,7 +159,9 @@ def _artifact_hf_uri( ) -def _dataset_uris_from_manifest(manifest) -> tuple[dict[str, str], dict[str, str]]: +def _dataset_uris_from_manifest( + manifest: CountryReleaseManifest, +) -> tuple[dict[str, str], dict[str, str]]: dataset_uris: dict[str, str] = {} dataset_repo_types: dict[str, str] = {} for name, reference in manifest.datasets.items(): @@ -171,7 +180,7 @@ def _dataset_uris_from_manifest(manifest) -> tuple[dict[str, str], dict[str, str def _dataset_uris_from_release( *, data_release: Mapping, - fallback_data_package, + fallback_data_package: DataPackageVersion, ) -> tuple[dict[str, str], dict[str, str]]: from policyengine.provenance.manifest import build_hf_uri diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py index 71fcfbfdf..3fc6de14f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py @@ -22,13 +22,14 @@ aggregates drift ~1e-7 relative. """ +from __future__ import annotations + import logging -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd from microdf import MicroDataFrame from policyengine.core import Simulation -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_executor.simulation_microdata import ( rebuild_entity_frame, @@ -37,6 +38,12 @@ SimulationOutputBuilder, ) +if TYPE_CHECKING: + from types import ModuleType + + from policyengine.tax_benefit_models.us.datasets import PolicyEngineUSDataset + from policyengine_observability import ObservabilityRuntime + logger = logging.getLogger(__name__) @@ -82,7 +89,9 @@ def concatenate_microdata( return sides -def _us_dataset_from_frames(frames: dict[str, pd.DataFrame], *, year: int): +def _us_dataset_from_frames( + frames: dict[str, pd.DataFrame], *, year: int +) -> PolicyEngineUSDataset: """Wrap concatenated entity tables as an in-memory US dataset. Weights need no special handling: they ride as ``_weight`` @@ -115,7 +124,7 @@ def build_national_output( *, country: str, simulation_params: dict[str, Any], - country_module, + country_module: ModuleType, year: int, resolved_data_version: str | None = None, runtime: ObservabilityRuntime | None = None, @@ -137,7 +146,7 @@ def build_national_output( for side, frames in sides.items() } - def stand_in(dataset) -> PrecomputedSimulation: + def stand_in(dataset: PolicyEngineUSDataset) -> PrecomputedSimulation: simulation = PrecomputedSimulation( dataset=dataset, tax_benefit_model_version=country_module.model, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py index fd76d9751..24bfb5547 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py @@ -15,7 +15,7 @@ class AppSettings(BaseSettings): @field_validator("environment", mode="before") @classmethod - def strip_environment(cls, v): + def strip_environment(cls, v: object) -> object: if isinstance(v, str): return v.strip() return v @@ -33,5 +33,5 @@ def strip_environment(cls, v): @lru_cache -def get_settings(): +def get_settings() -> AppSettings: return AppSettings() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py index 78d8b731d..4016de5ec 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py @@ -14,7 +14,7 @@ logger = logging.getLogger(__file__) -def create_router(runtime: ObservabilityRuntime): +def create_router(runtime: ObservabilityRuntime) -> APIRouter: router = APIRouter() @router.post("/simulate/economy/comparison", response_model=EconomyComparison) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py index 7f7b8370e..c51a28fc2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py @@ -13,15 +13,22 @@ bit-exact against a direct run. """ +from __future__ import annotations + import logging -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd +if TYPE_CHECKING: + from policyengine.core import Simulation + logger = logging.getLogger(__name__) -def extract_output_microdata(baseline, reform) -> dict[str, Any]: +def extract_output_microdata( + baseline: Simulation, reform: Simulation +) -> dict[str, Any]: """Columnar dump of each simulation's computed ``output_dataset.data`` per entity, for baseline and reform, plus per-column source dtypes. @@ -34,7 +41,7 @@ def extract_output_microdata(baseline, reform) -> dict[str, Any]: """ def dump( - simulation, + simulation: Simulation, ) -> tuple[dict[str, dict[str, list]], dict[str, dict[str, str]]]: columns: dict[str, dict[str, list]] = {} dtypes: dict[str, dict[str, str]] = {} diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py index e12028e68..7d9a1e95e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING from policyengine_simulation_executor.simulation_macro_output import ( BudgetaryImpact, @@ -17,8 +17,11 @@ _sum_output_variable, ) +if TYPE_CHECKING: + from policyengine.core import Simulation -def build_detailed_budget(analysis: Any) -> DetailedBudgetOutput: + +def build_detailed_budget(analysis: object) -> DetailedBudgetOutput: collection = getattr(analysis, "program_statistics", None) if isinstance(collection, DetailedBudgetOutput): return collection @@ -37,7 +40,9 @@ def build_detailed_budget(analysis: Any) -> DetailedBudgetOutput: return DetailedBudgetOutput(detailed_budget) -def build_budgetary_impact(country: str, baseline, reform) -> BudgetaryImpact: +def build_budgetary_impact( + country: str, baseline: Simulation, reform: Simulation +) -> BudgetaryImpact: tax_revenue_impact = _change_output_variable( baseline, reform, "household_tax", entity="household" ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index ee368501b..3207f0fa9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -2,7 +2,7 @@ from __future__ import annotations -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -60,11 +60,11 @@ class SimulationOutputBuilder: def __post_init__(self) -> None: self.country = self.country.lower() - def _span(self, name: str): + def _span(self, name: str) -> AbstractContextManager[object]: return self.runtime.span(name) if self.runtime is not None else nullcontext() @property - def analysis(self) -> Any: + def analysis(self) -> object: if self._analysis is None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.ECONOMIC_IMPACT_ANALYSIS)): self._analysis = self.country_module.economic_impact_analysis( @@ -148,14 +148,14 @@ def _build_intra_decile_output(self) -> IntraDecileOutput: self.baseline, self.reform ) - def _build_wealth_decile(self, wealth_decile: Any) -> DecileOutput | None: + def _build_wealth_decile(self, wealth_decile: object) -> DecileOutput | None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_WEALTH_DECILE)): return simulation_output_distribution.build_wealth_decile( self.country, wealth_decile ) def _build_intra_wealth_decile( - self, intra_wealth_decile: Any + self, intra_wealth_decile: object ) -> IntraDecileOutput | None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_INTRA_WEALTH_DECILE)): return simulation_output_distribution.build_intra_wealth_decile( @@ -171,15 +171,15 @@ def _build_cliff_impact(self) -> CliffImpactOutput | None: return simulation_output_cliff.build_cliff_impact(self.analysis) def _build_geographic_impact_output( - self, value: Any + self, value: object ) -> GeographicImpactOutput | None: return simulation_output_geographic.build_geographic_impact_output(value) - def _build_decile_output(self, collection: Any) -> DecileOutput: + def _build_decile_output(self, collection: object) -> DecileOutput: return simulation_output_distribution.build_decile_output(collection) def _build_intra_decile_output_from_collection( - self, collection: Any + self, collection: object ) -> IntraDecileOutput: return simulation_output_distribution.build_intra_decile_output_from_collection( collection @@ -188,10 +188,10 @@ def _build_intra_decile_output_from_collection( def _build_poverty_output( self, *, - baseline: Any, - reform: Any, - baseline_by_age: Any, - reform_by_age: Any, + baseline: object, + reform: object, + baseline_by_age: object, + reform_by_age: object, ) -> PovertyOutput: return simulation_output_poverty.build_poverty_output( country=self.country, @@ -204,8 +204,8 @@ def _build_poverty_output( def _build_poverty_by_gender_output( self, *, - baseline_by_gender: Any, - reform_by_gender: Any, + baseline_by_gender: object, + reform_by_gender: object, ) -> PovertyByGenderOutput: return simulation_output_poverty.build_poverty_by_gender_output( country=self.country, @@ -216,8 +216,8 @@ def _build_poverty_by_gender_output( def _build_poverty_by_race_output( self, *, - baseline_by_race: Any, - reform_by_race: Any, + baseline_by_race: object, + reform_by_race: object, ) -> PovertyByRaceOutput: return simulation_output_poverty.build_poverty_by_race_output( baseline_by_race=baseline_by_race, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py index 0e2953b87..746d1b8f2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py @@ -3,7 +3,6 @@ from __future__ import annotations from collections.abc import Mapping -from typing import Any from policyengine_simulation_executor.simulation_macro_output import ( CliffImpactInSimulation, @@ -12,7 +11,7 @@ from policyengine_simulation_executor.simulation_output_common import _output_model_dump -def build_cliff_impact(analysis: Any) -> CliffImpactOutput | None: +def build_cliff_impact(analysis: object) -> CliffImpactOutput | None: cliff_impact = getattr(analysis, "cliff_impact", None) if isinstance(cliff_impact, CliffImpactOutput): return cliff_impact diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py index 941b4ab72..e2fd0d7f2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py @@ -4,16 +4,21 @@ import logging import math -from collections.abc import Mapping +from collections.abc import Callable, Mapping from importlib import import_module -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import pandas as pd + from _typeshed import ConvertibleToFloat + from policyengine.core import Simulation import numpy as np logger = logging.getLogger(__name__) -def _number(value: Any, default: float = 0.0) -> float: +def _number(value: ConvertibleToFloat | None, default: float = 0.0) -> float: if value is None: return default try: @@ -25,7 +30,7 @@ def _number(value: Any, default: float = 0.0) -> float: return result -def _collection_records(collection: Any) -> list[dict[str, Any]]: +def _collection_records(collection: object) -> list[dict[str, Any]]: if collection is None: return [] dataframe = getattr(collection, "dataframe", None) @@ -36,7 +41,7 @@ def _collection_records(collection: Any) -> list[dict[str, Any]]: return [] -def _output_model_dump(value: Any) -> Any: +def _output_model_dump(value: object) -> object: if value is None: return None if hasattr(value, "model_dump"): @@ -46,13 +51,13 @@ def _output_model_dump(value: Any) -> Any: return None -def _entity_data(simulation, entity: str): +def _entity_data(simulation: Simulation, entity: str) -> pd.DataFrame: if simulation.output_dataset is None or simulation.output_dataset.data is None: simulation.ensure() return getattr(simulation.output_dataset.data, entity) -def _count_output_entities(simulation, entity: str, weight: str) -> float: +def _count_output_entities(simulation: Simulation, entity: str, weight: str) -> float: """The weighted number of entities: the plain sum of their weights. Output frames are microdf ``MicroDataFrame``s weighted by the entity's @@ -66,7 +71,7 @@ def _count_output_entities(simulation, entity: str, weight: str) -> float: return float(np.asarray(data[weight], dtype=float).sum()) -def _sum_output_variable(simulation, variable: str, entity: str) -> float: +def _sum_output_variable(simulation: Simulation, variable: str, entity: str) -> float: data = _entity_data(simulation, entity) if variable in data.columns: return float(data[variable].sum()) @@ -83,7 +88,9 @@ def _sum_output_variable(simulation, variable: str, entity: str) -> float: return float(output.result) -def _change_output_variable(baseline, reform, variable: str, entity: str) -> float: +def _change_output_variable( + baseline: Simulation, reform: Simulation, variable: str, entity: str +) -> float: baseline_data = _entity_data(baseline, entity) reform_data = _entity_data(reform, entity) if variable in baseline_data.columns and variable in reform_data.columns: @@ -102,16 +109,18 @@ def _change_output_variable(baseline, reform, variable: str, entity: str) -> flo return float(output.result) -def _output_module_function(module_name: str, name: str): +def _output_module_function(module_name: str, name: str) -> Callable[..., object]: module = import_module(f"policyengine.outputs.{module_name}") return getattr(module, name) -def _poverty_module_function(name: str): +def _poverty_module_function(name: str) -> Callable[..., object]: return _output_module_function("poverty", name) -def _try_compute_output(label: str, fn, *args, **kwargs): +def _try_compute_output[**P, T]( + label: str, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs +) -> T | None: try: return fn(*args, **kwargs) except Exception as exc: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py index 59be0d9af..32032e308 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING from policyengine_simulation_executor.simulation_macro_output import ( DecileOutput, @@ -14,6 +14,9 @@ _try_compute_output, ) +if TYPE_CHECKING: + from policyengine.core import Simulation + INTRA_DECILE_COLUMNS = { "Lose more than 5%": "lose_more_than_5pct", "Lose less than 5%": "lose_less_than_5pct", @@ -23,11 +26,11 @@ } -def build_decile(analysis: Any) -> DecileOutput: +def build_decile(analysis: object) -> DecileOutput: return build_decile_output(getattr(analysis, "decile_impacts", None)) -def build_decile_output(collection: Any) -> DecileOutput: +def build_decile_output(collection: object) -> DecileOutput: if isinstance(collection, DecileOutput): return collection average: dict[str, float] = {} @@ -45,7 +48,9 @@ def build_decile_output(collection: Any) -> DecileOutput: return DecileOutput(average=average, relative=relative) -def build_intra_decile_output(baseline, reform) -> IntraDecileOutput: +def build_intra_decile_output( + baseline: Simulation, reform: Simulation +) -> IntraDecileOutput: from policyengine.outputs.intra_decile_impact import ( compute_intra_decile_impacts, ) @@ -61,21 +66,23 @@ def build_intra_decile_output(baseline, reform) -> IntraDecileOutput: return build_intra_decile_output_from_collection(collection) -def build_wealth_decile(country: str, wealth_decile: Any) -> DecileOutput | None: +def build_wealth_decile(country: str, wealth_decile: object) -> DecileOutput | None: if country != "uk": return None return build_decile_output(wealth_decile) def build_intra_wealth_decile( - country: str, intra_wealth_decile: Any + country: str, intra_wealth_decile: object ) -> IntraDecileOutput | None: if country != "uk": return None return build_intra_decile_output_from_collection(intra_wealth_decile) -def build_intra_decile_output_from_collection(collection: Any) -> IntraDecileOutput: +def build_intra_decile_output_from_collection( + collection: object, +) -> IntraDecileOutput: if isinstance(collection, IntraDecileOutput): return collection deciles: dict[str, list[float]] = {label: [] for label in INTRA_DECILE_COLUMNS} diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py index 4deb42a05..2bf92c1b9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py @@ -24,6 +24,7 @@ ) if TYPE_CHECKING: + from policyengine.core import Simulation from policyengine.data.uk_geography_assets import UKGeographyAssetSpec from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata @@ -89,7 +90,7 @@ def _policyengine_us_district_metadata() -> tuple[dict[int, str], frozenset[str] ) -def build_geographic_impact_output(value: Any) -> GeographicImpactOutput | None: +def build_geographic_impact_output(value: object) -> GeographicImpactOutput | None: if isinstance(value, GeographicImpactOutput): return value records = _output_model_dump(value) @@ -105,7 +106,7 @@ def build_geographic_impact_output(value: Any) -> GeographicImpactOutput | None: def build_congressional_district_impact_output( - value: Any, + value: object, ) -> CongressionalDistrictImpactOutput | None: if isinstance(value, CongressionalDistrictImpactOutput): return value @@ -187,7 +188,11 @@ def _should_build_us_congressional_district_impact(region_code: str | None) -> b def build_congressional_district_impact( - country: str, baseline, reform, *, region_code: str | None = None + country: str, + baseline: Simulation, + reform: Simulation, + *, + region_code: str | None = None, ) -> CongressionalDistrictImpactOutput | None: if country != "us": return None @@ -207,7 +212,7 @@ def compute_and_format() -> CongressionalDistrictImpactOutput | None: return _try_compute_output("congressional district impacts", compute_and_format) -def _output_household(simulation) -> pd.DataFrame: +def _output_household(simulation: Simulation) -> pd.DataFrame: return pd.DataFrame(simulation.output_dataset.data.household) @@ -230,7 +235,7 @@ def _require_uk_area_column( def build_uk_constituency_impact( - country: str, baseline, reform + country: str, baseline: Simulation, reform: Simulation ) -> GeographicImpactOutput | None: if country != "uk": return None @@ -256,8 +261,8 @@ def build_uk_constituency_impact( def build_uk_local_authority_impact( country: str, - baseline, - reform, + baseline: Simulation, + reform: Simulation, *, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> GeographicImpactOutput | None: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py index a9d8639ce..7bd1366a5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py @@ -2,8 +2,6 @@ from __future__ import annotations -from typing import Any - from policyengine_simulation_executor.simulation_macro_output import ( BaselineReformValue, InequalityOutput, @@ -11,7 +9,7 @@ from policyengine_simulation_executor.simulation_output_common import _number -def build_inequality(analysis: Any) -> InequalityOutput: +def build_inequality(analysis: object) -> InequalityOutput: baseline = getattr(analysis, "baseline_inequality", None) reform = getattr(analysis, "reform_inequality", None) if isinstance(baseline, InequalityOutput): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py index 447ba3e1f..7033fe5ab 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py @@ -2,15 +2,13 @@ from __future__ import annotations -from typing import Any - from policyengine_simulation_executor.simulation_macro_output import ( LaborSupplyResponseOutput, ) from policyengine_simulation_executor.simulation_output_common import _output_model_dump -def build_labor_supply_response(analysis: Any) -> LaborSupplyResponseOutput | None: +def build_labor_supply_response(analysis: object) -> LaborSupplyResponseOutput | None: labor_supply_response = getattr(analysis, "labor_supply_response", None) if isinstance(labor_supply_response, LaborSupplyResponseOutput): return labor_supply_response diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py index de63db21d..fb260ce52 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py @@ -24,6 +24,8 @@ if TYPE_CHECKING: from collections.abc import Iterable, Mapping + from policyengine.core import Simulation + US_POVERTY_TYPES = { "spm": "poverty", "spm_deep": "deep_poverty", @@ -110,7 +112,9 @@ def _race_poverty_output(values: dict[str, dict[str, float]]) -> RacePovertyOutp ) -def build_poverty_outputs(country: str, baseline, reform, analysis: Any): +def build_poverty_outputs( + country: str, baseline: Simulation, reform: Simulation, analysis: object +) -> PovertyModuleOutputs: prefix = "us" if country == "us" else "uk" baseline_poverty_by_age = _try_compute_output( "baseline poverty by age", @@ -172,10 +176,10 @@ def build_poverty_outputs(country: str, baseline, reform, analysis: Any): def build_poverty_output( *, country: str, - baseline: Any, - reform: Any, - baseline_by_age: Any, - reform_by_age: Any, + baseline: object, + reform: object, + baseline_by_age: object, + reform_by_age: object, ) -> PovertyOutput: if isinstance(baseline, PovertyOutput): return baseline @@ -206,8 +210,8 @@ def build_poverty_output( def build_poverty_by_gender_output( *, country: str, - baseline_by_gender: Any, - reform_by_gender: Any, + baseline_by_gender: object, + reform_by_gender: object, ) -> PovertyByGenderOutput: if isinstance(baseline_by_gender, PovertyByGenderOutput): return baseline_by_gender @@ -230,8 +234,8 @@ def build_poverty_by_gender_output( def build_poverty_by_race_output( *, - baseline_by_race: Any, - reform_by_race: Any, + baseline_by_race: object, + reform_by_race: object, ) -> PovertyByRaceOutput: if isinstance(baseline_by_race, PovertyByRaceOutput): return baseline_by_race diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index 6004f9283..871656ae5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -6,6 +6,8 @@ No Modal dependencies here. """ +from __future__ import annotations + import contextlib import hashlib import json @@ -13,13 +15,11 @@ import os import tempfile import time -from collections.abc import Iterator from dataclasses import dataclass from datetime import datetime from importlib import import_module -from typing import Any +from typing import TYPE_CHECKING, Any -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri from policyengine_simulation_observability.stages import ( ANNUAL_IMPACT_STAGES, @@ -35,6 +35,16 @@ SimulationOutputBuilder, ) +if TYPE_CHECKING: + from collections.abc import Iterator + from types import ModuleType + + from policyengine.core import Dataset, Policy, Region, Simulation + from policyengine.core.scoping_strategy import ScopingStrategy + from policyengine.provenance.manifest import ArtifactPathReference + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_observability.telemetry import TelemetryEnvelope + logger = logging.getLogger(__name__) os.environ.setdefault("POLICYENGINE_SKIP_COUNTRY_IMPORTS", "1") @@ -182,7 +192,7 @@ def _parse_year(params: dict[str, Any]) -> int: return int(value) -def _parse_policy_period(period_key: Any) -> tuple[datetime, datetime | None]: +def _parse_policy_period(period_key: object) -> tuple[datetime, datetime | None]: """Parse effective dates or inclusive ``start.stop`` policy periods.""" text = str(period_key) parts = text.split(".") @@ -201,7 +211,7 @@ def _parse_policy_period(period_key: Any) -> tuple[datetime, datetime | None]: return start, stop -def _normalise_period_key(period_key: Any) -> str: +def _normalise_period_key(period_key: object) -> str: """Validate policy period keys without discarding their stop dates.""" text = str(period_key) _parse_policy_period(text) @@ -224,7 +234,9 @@ def _normalise_policy(policy: dict[str, Any] | None) -> dict[str, Any] | None: return normalised -def _compile_interval_policy(policy, *, country_module, year: int): +def _compile_interval_policy( + policy: dict[str, Any] | None, *, country_module: ModuleType, year: int +) -> dict[str, Any] | Policy | None: """Use public wrapper values for intervals its dict compiler cannot parse.""" if not policy or not any( isinstance(value, dict) and any("." in str(key) for key in value) @@ -260,7 +272,9 @@ def _compile_interval_policy(policy, *, country_module, year: int): return Policy(name="Reform", parameter_values=parameter_values) -def _runtime_policy(policy, *, country_module, year: int): +def _runtime_policy( + policy: dict[str, Any] | None, *, country_module: ModuleType, year: int +) -> dict[str, Any] | Policy | None: """A request's policy as the wrapper must receive it. Normalises the period keys, then compiles a policy with interval keys @@ -276,7 +290,7 @@ def _resolve_dataset_reference(country: str) -> str: return resolve_bundle_dataset_name(country, None) -def _normalise_region_code(country: str, region: Any) -> str: +def _normalise_region_code(country: str, region: object) -> str: if region is None or str(region).strip() == "": return country @@ -305,7 +319,7 @@ def _normalise_region_code(country: str, region: Any) -> str: return f"{prefix}/{value}" -def _build_uk_weight_replacement_region(region_code: str): +def _build_uk_weight_replacement_region(region_code: str) -> Region | None: if "/" not in region_code: return None @@ -345,7 +359,9 @@ def _build_uk_weight_replacement_region(region_code: str): def _require_certified_uk_weight_matrix( - scoping_strategy, dataset, dataset_selection + scoping_strategy: ScopingStrategy | None, + dataset: Dataset, + dataset_selection: DatasetSelection, ) -> None: """Bind a weight-matrix region to the release of the dataset it reweights. @@ -385,7 +401,7 @@ def _require_certified_uk_weight_matrix( manifest = get_release_manifest("uk") package = manifest.data_package - def release(reference) -> tuple[str, str]: + def release(reference: ArtifactPathReference) -> tuple[str, str]: return ( reference.repo_id or package.repo_id, reference.revision or package.release_manifest_revision or package.version, @@ -454,9 +470,9 @@ def release(reference) -> tuple[str, str]: def _region_parent_dataset_reference( - country_module, + country_module: ModuleType, country: str, - region, + region: Region, ) -> str: parent_code = getattr(region, "parent_code", None) visited: set[str] = set() @@ -481,7 +497,7 @@ def _region_parent_dataset_reference( return _resolve_dataset_reference(country) -def _reject_unscoped_us_place_region(region_code: str, region) -> None: +def _reject_unscoped_us_place_region(region_code: str, region: Region) -> None: if not region_code.startswith("place/"): return if getattr(region, "dataset_path", None) is not None: @@ -496,7 +512,7 @@ def _reject_unscoped_us_place_region(region_code: str, region) -> None: def _resolve_region( *, - country_module, + country_module: ModuleType, country: str, params: dict[str, Any], ) -> RegionResolution: @@ -544,7 +560,7 @@ def _resolve_region( def _resolve_region_group( *, - country_module, + country_module: ModuleType, country: str, params: dict[str, Any], ) -> RegionResolution: @@ -588,7 +604,7 @@ def _resolve_region_group( ) -def _country_module(country: str): +def _country_module(country: str) -> ModuleType: country = country.lower() if country not in {"us", "uk"}: raise ValueError(f"Unsupported country: {country}") @@ -649,8 +665,8 @@ def _load_dataset( params: dict[str, Any], *, selection: DatasetSelection, - country_module=None, -): + country_module: ModuleType | None = None, +) -> Dataset: country = params.get("country", "us").lower() year = _parse_year(params) country_module = country_module or _country_module(country) @@ -678,12 +694,12 @@ def _load_dataset( def _build_simulation( params: dict[str, Any], *, - dataset, + dataset: Dataset, dataset_selection: DatasetSelection, policy: dict[str, Any] | None, - scoping_strategy=None, + scoping_strategy: ScopingStrategy | None = None, region_code: str | None = None, -): +) -> Simulation: from policyengine.core import Simulation from policyengine_simulation_executor.baseline_artifacts import ( @@ -694,7 +710,7 @@ def _build_simulation( selection = normalize_runtime_spm(params) params = {**params, **({"spm": selection} if selection is not None else {})} - spm_kwargs = {"spm": selection} if selection is not None else {} + spm_kwargs: dict[str, Any] = {"spm": selection} if selection is not None else {} country = params.get("country", "us") country_module = _country_module(country) policy = _compile_interval_policy( @@ -862,7 +878,7 @@ def _set_runtime_attributes( *, runtime: ObservabilityRuntime, simulation_params: dict[str, Any], - telemetry, + telemetry: TelemetryEnvelope | None, metadata: dict[str, Any], country: str, ) -> None: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index d76ad3fcb..8cb5696f8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -1,8 +1,11 @@ """Certified SPM runtime selection, independent of caller-supplied metadata.""" +from __future__ import annotations + import json from collections.abc import Mapping from functools import lru_cache +from typing import TYPE_CHECKING, Any from policyengine_simulation_contract.spm import ( SPMCapability, @@ -17,16 +20,19 @@ validate_spm_calculation_provenance, ) +if TYPE_CHECKING: + from spm_calculator.rolling_forecast import SPMForecast + @lru_cache(maxsize=4) -def _forecast(expected_sha256): +def _forecast(expected_sha256: str | None) -> SPMForecast: from spm_calculator import load_forecast return load_forecast(expected_sha256=expected_sha256) @lru_cache(maxsize=1) -def runtime_spm_capability(): +def runtime_spm_capability() -> SPMCapability | None: """The installed image's certified SPM capability, or None. Memoized for the life of the process because every input is a property @@ -83,7 +89,9 @@ def runtime_spm_capability(): @lru_cache(maxsize=64) -def _prevalidate_selection(selection_json: str, start_year: int, window_size: int): +def _prevalidate_selection( + selection_json: str, start_year: int, window_size: int +) -> None: """Prove a resolved selection can be measured, before anything expensive. Keyed on the resolved selection and the year range, which is everything @@ -139,7 +147,7 @@ def _prevalidate_selection(selection_json: str, start_year: int, window_size: in ) -def reset_spm_runtime_caches(): +def reset_spm_runtime_caches() -> None: """Forget what was memoized from the installed environment. Production never needs this: a container's image is fixed for the life @@ -154,7 +162,7 @@ def reset_spm_runtime_caches(): clear() -def normalize_runtime_spm(params): +def normalize_runtime_spm(params: dict[str, Any]) -> dict[str, Any] | None: """Resolve before dataset loading, artifact lookup, or child submission.""" from policyengine_simulation_executor.release_bundle import ( get_country_release_bundle, @@ -277,7 +285,13 @@ def simulation_spm_calculation_provenance( ) from exc -def simulation_spm_result(baseline, reform, selection, *, expected_year=None): +def simulation_spm_result( + baseline: object, + reform: object, + selection: object, + *, + expected_year: int | None = None, +) -> dict[str, object]: if selection is None: return {} baseline_provenance = simulation_spm_calculation_provenance(baseline, selection) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py index 798c091af..6d88745ce 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py @@ -239,7 +239,7 @@ def deserialize_uk_local_authority_metadata( class Stage12ArtifactStore: - def __init__(self, bucket_name: str, *, store: ArtifactStore | None = None): + def __init__(self, bucket_name: str, *, store: ArtifactStore | None = None) -> None: if not bucket_name: raise ValueError( "Stage 12 requires an explicitly configured private bucket" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index b61a3c0f1..4c8b05993 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -15,6 +15,8 @@ from collections.abc import Mapping import pandas as pd + from policyengine.tax_benefit_models.uk.datasets import PolicyEngineUKDataset + from policyengine.tax_benefit_models.us.datasets import PolicyEngineUSDataset from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from pydantic import JsonValue @@ -86,7 +88,7 @@ def _dataset_from_frames( *, country: CountryId, year: int, -): +) -> PolicyEngineUSDataset | PolicyEngineUKDataset: from microdf import MicroDataFrame fields = {} @@ -159,7 +161,10 @@ def build_aggregate_report( ), } - def stand_in(dataset, policy: dict[str, JsonValue]): + def stand_in( + dataset: PolicyEngineUSDataset | PolicyEngineUKDataset, + policy: dict[str, JsonValue], + ) -> PrecomputedSimulation: simulation = PrecomputedSimulation( dataset=dataset, tax_benefit_model_version=country_module.model, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 3063ddbd3..1e6aa584a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -60,6 +60,7 @@ if TYPE_CHECKING: from collections.abc import Callable, Mapping + from contextlib import AbstractContextManager from policyengine_observability import ObservabilityRuntime @@ -280,7 +281,7 @@ def coordinate_report( def stage_scope( stage: Stage, attributes: Mapping[str, object] | None = None, - ): + ) -> AbstractContextManager[object]: if runtime is None: return nullcontext() return runtime.operation( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py index 8f97f2237..e30e3c6f2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py @@ -10,6 +10,7 @@ from functools import lru_cache from importlib.resources import files from types import MappingProxyType +from typing import TYPE_CHECKING import pandas as pd from policyengine_simulation_contract.uk_geography import ( @@ -17,6 +18,9 @@ UKLocalAuthorityMetadata, ) +if TYPE_CHECKING: + from importlib.resources.abc import Traversable + _CODE_PATTERN = re.compile(r"^[A-Z]\d{8}$") _RESOURCE_PACKAGE = "policyengine_simulation_executor" _RESOURCE_DIRECTORY = ("static_runtime_files", "uk_local_authorities") @@ -116,7 +120,7 @@ def metadata_for( ) -def _resource_path(filename: str): +def _resource_path(filename: str) -> Traversable: path = files(_RESOURCE_PACKAGE) for part in _RESOURCE_DIRECTORY: path = path.joinpath(part) diff --git a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py index f5ac439a0..3866aff93 100644 --- a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py +++ b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py @@ -1,6 +1,8 @@ """Fixtures for gateway endpoint tests.""" +from collections.abc import ItemsView, KeysView from copy import deepcopy +from typing import Self import pytest from policyengine_simulation_contract.macro_output import ( @@ -159,28 +161,28 @@ def resolve_test_dataset_uri( class MockDict: """Mock for Modal.Dict to simulate version registry.""" - def __init__(self, data: dict): + def __init__(self, data: dict) -> None: self._data = data - def __getitem__(self, key: str): + def __getitem__(self, key: str) -> object: if key not in self._data: raise KeyError(key) return self._data[key] - def __setitem__(self, key: str, value): + def __setitem__(self, key: str, value: object) -> None: self._data[key] = value - def get(self, key: str, default=None): + def get(self, key: str, default: object = None) -> object: return self._data.get(key, default) - def keys(self): + def keys(self) -> KeysView[str]: return self._data.keys() - def items(self): + def items(self) -> ItemsView[str, object]: return self._data.items() @classmethod - def from_name(cls, name: str): + def from_name(cls, name: str) -> Self: """Mock from_name that returns a MockDict based on name.""" raise NotImplementedError("Mock not configured") @@ -191,7 +193,7 @@ class MockFunctionCall: registry = {} from_id_errors = {} - def __init__(self, object_id: str = "mock-job-id-123"): + def __init__(self, object_id: str = "mock-job-id-123") -> None: self.object_id = object_id value = BaselineReformValue(baseline=0.0, reform=0.0) age_poverty = AgePovertyOutput( @@ -241,7 +243,7 @@ def __init__(self, object_id: str = "mock-job-id-123"): self.running = False self.__class__.registry[object_id] = self - def get(self, timeout: int = 0): + def get(self, timeout: int = 0) -> object: if self.running: raise TimeoutError if self.error is not None: @@ -249,7 +251,7 @@ def get(self, timeout: int = 0): return self.result @classmethod - def from_id(cls, object_id: str): + def from_id(cls, object_id: str) -> "MockFunctionCall": if object_id in cls.from_id_errors: raise cls.from_id_errors[object_id] if object_id not in cls.registry: @@ -260,7 +262,7 @@ def from_id(cls, object_id: str): class MockFunction: """Mock for Modal Function.""" - def __init__(self): + def __init__(self) -> None: self.last_payload = None self.last_observability_context = None self.last_from_name_call = None @@ -279,7 +281,7 @@ def call_for(self, object_id: str) -> MockFunctionCall: class BoundMockFunction: """Function handle returned by Modal.Function.from_name.""" - def __init__(self, recorder: MockFunction, app_name: str, func_name: str): + def __init__(self, recorder: MockFunction, app_name: str, func_name: str) -> None: self.recorder = recorder self.app_name = app_name self.func_name = func_name @@ -318,7 +320,7 @@ class OutputExpiredError(Exception): @pytest.fixture -def mock_modal(monkeypatch): +def mock_modal(monkeypatch: pytest.MonkeyPatch) -> dict[str, object]: """Patch Modal calls in the gateway endpoints module.""" from policyengine_simulation_contract import budget_window_state, dataset_uri from policyengine_simulation_gateway import endpoints @@ -333,7 +335,7 @@ def mock_modal(monkeypatch): class MockModalDict: @staticmethod - def from_name(name: str, create_if_missing: bool = False): + def from_name(name: str, create_if_missing: bool = False) -> MockDict: if create_if_missing and name not in mock_dicts: mock_dicts[name] = {} if name not in mock_dicts: @@ -342,7 +344,7 @@ def from_name(name: str, create_if_missing: bool = False): class MockModalFunction: @staticmethod - def from_name(app_name: str, func_name: str): + def from_name(app_name: str, func_name: str) -> BoundMockFunction: mock_func.last_from_name_call = (app_name, func_name) return mock_func.bind(app_name, func_name) diff --git a/projects/policyengine-simulation-gateway/fixtures/package_imports.py b/projects/policyengine-simulation-gateway/fixtures/package_imports.py index b6769a469..d508167d1 100644 --- a/projects/policyengine-simulation-gateway/fixtures/package_imports.py +++ b/projects/policyengine-simulation-gateway/fixtures/package_imports.py @@ -10,7 +10,8 @@ import pytest if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Callable, Iterator + from types import ModuleType GATEWAY_MODEL_MODULE = "policyengine_simulation_contract.gateway_models" GATEWAY_ENDPOINTS_MODULE = "policyengine_simulation_gateway.endpoints" @@ -61,10 +62,12 @@ def isolated_gateway_model_import_modules() -> Iterator[None]: @pytest.fixture() -def import_gateway_models(isolated_gateway_model_import_modules): +def import_gateway_models( + isolated_gateway_model_import_modules: None, +) -> Callable[[], ModuleType]: """Import gateway models from a clean module state.""" - def import_models(): + def import_models() -> ModuleType: return importlib.import_module(GATEWAY_MODEL_MODULE) return import_models diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py index 2b4535033..ee05f39a1 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py @@ -8,14 +8,20 @@ and spawns jobs on those apps. """ +from __future__ import annotations + import os from pathlib import Path +from typing import TYPE_CHECKING import modal from policyengine_simulation_observability.observability import ( modal_image_environment, ) +if TYPE_CHECKING: + from fastapi import FastAPI + # Stable app name - this should rarely change app = modal.App("policyengine-simulation-gateway") gateway_auth_secret = modal.Secret.from_name("policyengine-gateway-auth") @@ -64,7 +70,7 @@ def build_gateway_image() -> modal.Image: @app.function(image=gateway_image, secrets=[gateway_auth_secret]) @modal.asgi_app() -def web_app(): +def web_app() -> FastAPI: """ FastAPI gateway for simulation job submission and polling. diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py index 620e7f774..c5e30a5c4 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py @@ -10,14 +10,17 @@ ) if TYPE_CHECKING: - from fastapi import FastAPI, Request + from fastapi import FastAPI, Request, Response + from starlette.middleware.base import RequestResponseEndpoint def install_observability_id_middleware(app: FastAPI) -> None: """Store an incoming identifier candidate without creating a workflow.""" @app.middleware("http") - async def observability_id_transport(request: Request, call_next): + async def observability_id_transport( + request: Request, call_next: RequestResponseEndpoint + ) -> Response: request.state.incoming_observability_id = normalize_observability_id( request.headers.get(OBSERVABILITY_ID_HEADER) ) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py index d23832d98..6fd6796ac 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -5,10 +5,11 @@ import logging import re from dataclasses import dataclass -from typing import TypedDict +from typing import Any, TypedDict import modal from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi.responses import JSONResponse from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.budget_window_state import ( build_batch_status_response, @@ -23,6 +24,7 @@ BudgetWindowBatchRequest, BudgetWindowBatchStatusResponse, BudgetWindowBatchSubmitResponse, + GatewayRequestBase, HealthResponse, JobStatusResponse, JobSubmitResponse, @@ -36,6 +38,7 @@ from policyengine_simulation_contract.hf_dataset import ( HuggingFaceDatasetReferenceError, ) +from policyengine_simulation_contract.json_types import JsonObject from policyengine_simulation_contract.spm import ( SPMCapability, SPMInputError, @@ -91,7 +94,7 @@ class RouteResolution: country_model_version: str | None = None -def _job_metadata_store(): +def _job_metadata_store() -> "modal.Dict": return modal.Dict.from_name(JOB_METADATA_DICT_NAME, create_if_missing=True) @@ -158,7 +161,7 @@ def _resolve_dataset_uri_from_app_bundle( ) -def _modal_exception_class(name: str): +def _modal_exception_class(name: str) -> type[BaseException] | None: exception_module = getattr(modal, "exception", None) if exception_module is None: return None @@ -176,7 +179,7 @@ def _is_modal_job_not_found(exc: BaseException) -> bool: ) -def _optional_modal_dict(name: str): +def _optional_modal_dict(name: str) -> "modal.Dict | None": try: return modal.Dict.from_name(name) except KeyError: @@ -629,7 +632,11 @@ def _certified_model_version(country: str, route: RouteResolution) -> str | None return route.country_model_version -def _resolve_request_spm(request, bundle, route): +def _resolve_request_spm( + request: GatewayRequestBase, + bundle: PolicyEngineBundle, + route: RouteResolution, +) -> JsonObject | None: selection = resolve_spm_selection( request.country, request.spm, @@ -645,7 +652,10 @@ def _resolve_request_spm(request, bundle, route): return selection -def _bundle_payload(bundle: PolicyEngineBundle, **dump_kwargs) -> dict: +def _bundle_payload( + bundle: PolicyEngineBundle, + **dump_kwargs: Any, # noqa: ANN401 # forwarded unchanged to BaseModel.model_dump +) -> dict: """Dump a bundle, omitting fields added after the legacy body when unset. The 202 and 500 job bodies splat this dict in directly, bypassing the @@ -778,7 +788,7 @@ async def submit_simulation( request: SimulationRequest, http_request: Request, response: Response, -): +) -> JobSubmitResponse | JSONResponse: """ Submit a simulation job. @@ -888,7 +898,7 @@ async def submit_budget_window_batch( request: BudgetWindowBatchRequest, http_request: Request, response: Response, -): +) -> BudgetWindowBatchSubmitResponse | JSONResponse: """ Submit a budget-window batch job. """ @@ -981,7 +991,9 @@ async def submit_budget_window_batch( response_model_exclude_none=True, dependencies=[Depends(require_auth)], ) -async def get_job_status(job_id: str, request: Request, response: Response): +async def get_job_status( + job_id: str, request: Request, response: Response +) -> JobStatusResponse | JSONResponse: """ Poll for job status. @@ -1064,7 +1076,7 @@ async def get_budget_window_job_status( batch_job_id: str, request: Request, response: Response, -): +) -> JSONResponse: """ Poll for budget-window batch status. """ diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py index 7fa67f982..207e5d76d 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py @@ -137,7 +137,7 @@ async def ping(request: PingRequest) -> PingResponse: return app -def main(): +def main() -> None: """Generate OpenAPI spec and write to artifacts.""" app = create_openapi_app() openapi_spec = app.openapi() diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py index c658c0e32..79428a456 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py @@ -15,7 +15,7 @@ class ServerErrorResponse(JSONResponse): """Shared 500 JSON response.""" - def __init__(self, content: dict): + def __init__(self, content: dict) -> None: super().__init__(status_code=500, content=content) @@ -38,7 +38,7 @@ def batch_status_response( response: BudgetWindowBatchStatusResponse, *, headers: dict[str, str] | None = None, -): +) -> JSONResponse: payload = batch_status_payload(response) if response.status in {"submitted", "running"}: return JSONResponse(status_code=202, content=payload, headers=headers) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py index aecb829c2..b0651d821 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py @@ -62,7 +62,7 @@ def smoke_import_gateway() -> dict: @app.local_entrypoint() -def main(): +def main() -> None: report = smoke_import_gateway.remote() print(report) assert report["modules_imported"] > 0 diff --git a/ruff.toml b/ruff.toml index a22dc8cc2..2aa803cb6 100644 --- a/ruff.toml +++ b/ruff.toml @@ -13,8 +13,15 @@ select = [ "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W", # Group 2: typing-only imports. "TC", + # Group 3: annotations. + "ANN", ] +[lint.per-file-ignores] +# 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"] + [lint.flake8-type-checking] # Classes whose annotations Pydantic evaluates at runtime, so their imports must # stay at runtime. Ruff checks direct bases only, so every project model base diff --git a/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py b/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py index f6b4b7262..8e6ea09d6 100644 --- a/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py +++ b/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py @@ -1,7 +1,14 @@ """Proof fixture: a Stage 12 runtime that creates its own table. The rule must reject it.""" +from __future__ import annotations + +from typing import TYPE_CHECKING + from sqlalchemy import text +if TYPE_CHECKING: + from sqlalchemy import Connection + -def ensure_comparison_table(connection) -> None: +def ensure_comparison_table(connection: Connection) -> None: connection.execute(text("CREATE TABLE stage12_evaluation_reports (id uuid)")) diff --git a/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py b/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py index 6dddf8588..b8cf2de0c 100644 --- a/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py +++ b/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py @@ -1,7 +1,14 @@ """Proof fixture: a Stage 12 runtime that writes rows with typed DML. The rule must accept it.""" +from __future__ import annotations + +from typing import TYPE_CHECKING + from sqlalchemy import insert +if TYPE_CHECKING: + from sqlalchemy import Connection, Table + -def record_report(connection, reports_table, report_id: str) -> None: +def record_report(connection: Connection, reports_table: Table, report_id: str) -> None: connection.execute(insert(reports_table).values(id=report_id)) diff --git a/rules/registry.toml b/rules/registry.toml index 10879a735..ae72bb9f1 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -710,7 +710,26 @@ 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" +invocation = "ruff check ." + +[[rule]] +id = "sim.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 = ["projects/*/tests/**/*.py", "libs/*/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" 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