From fda75385d0cd537f3adee91e3732297a447431b9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:39:18 +0400 Subject: [PATCH 01/10] Enable ruff group 2 (typing-only imports) and fix everything it finds Select TC and move imports used only in annotations under TYPE_CHECKING. Pydantic evaluates model annotations and FastAPI evaluates endpoint annotations at runtime, so ruff.toml lists the runtime-evaluated model base classes (every project base, since ruff only checks direct bases) and the FastAPI/APIRouter route decorators. scripts/rules matches policyengine-api. Co-Authored-By: Claude Opus 5.5 --- .../observability/emitters.py | 9 +++-- .../tests/test_stage12_execution.py | 12 ++++--- .../errors.py | 5 +-- .../observability.py | 6 ++-- .../stages.py | 5 ++- .../policyengine_stage12_persistence/store.py | 10 +++--- .../validation.py | 5 ++- .../validators.py | 5 ++- .../tests/simulation/test_calculate.py | 2 +- .../src/policyengine_simulation_entry/auth.py | 9 +++-- .../policyengine_simulation_entry/backend.py | 9 +++-- .../stage12_adapter.py | 10 +++--- .../stage12_backend.py | 9 +++-- .../tests/conftest.py | 5 ++- .../tests/test_app.py | 5 ++- .../tests/test_deploy_workflow_scripts.py | 5 ++- .../tests/test_openapi.py | 10 +++--- ...est_policyengine_package_update_scripts.py | 5 ++- .../src/modal/budget_window_batch.py | 6 ++-- .../src/modal/budget_window_scheduler.py | 6 ++-- .../src/modal/bundle_data.py | 5 ++- .../src/modal/fanout.py | 6 ++-- .../src/modal/segmented_national.py | 6 ++-- .../src/modal/utils/record_deployment.py | 5 ++- .../modal/utils/validate_v2_worker_release.py | 6 ++-- .../src/modal/v2_app.py | 5 ++- .../precompute.py | 5 ++- .../simulation_output_builder.py | 8 +++-- .../simulation_output_geographic.py | 2 +- .../simulation_output_poverty.py | 6 ++-- .../stage12_artifacts.py | 10 +++--- .../stage12_bundle/assertions.py | 9 +++-- .../stage12_bundle/manifest.py | 5 ++- .../stage12_result_comparison.py | 6 ++-- .../stage12_runtime/aggregation.py | 15 ++++---- .../stage12_runtime/comparison.py | 12 ++++--- .../stage12_runtime/coordination.py | 9 +++-- .../stage12_runtime/dependencies.py | 25 +++++++------ .../stage12_runtime/output_planning.py | 10 +++--- .../stage12_runtime/simulation.py | 12 ++++--- .../stage12_worker_validation.py | 7 ++-- ...est_policyengine_package_update_scripts.py | 5 ++- .../tests/test_stage12_output_planning.py | 7 ++-- .../fixtures/package_imports.py | 5 ++- .../correlation.py | 6 +++- .../responses.py | 10 ++++-- ruff.toml | 36 +++++++++++++++++++ scripts/rules/baseline.py | 5 ++- scripts/rules/discover.py | 5 ++- scripts/rules/model.py | 8 +++-- scripts/rules/resolve.py | 8 +++-- scripts/rules/validate.py | 6 ++-- 52 files changed, 287 insertions(+), 126 deletions(-) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py index 376aaf46d..3a06f307d 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py @@ -1,12 +1,15 @@ from __future__ import annotations -from collections.abc import Mapping from contextlib import AbstractContextManager from dataclasses import dataclass, field -from typing import Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol from .config import ObservabilityConfig -from .contracts import SimulationLifecycleEvent, TracerArtifactManifest + +if TYPE_CHECKING: + from collections.abc import Mapping + + from .contracts import SimulationLifecycleEvent, TracerArtifactManifest class NoOpSpan(AbstractContextManager["NoOpSpan"]): diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_execution.py b/libs/policyengine-simulation-contract/tests/test_stage12_execution.py index 24866e3fe..dda655d46 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_execution.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_execution.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import cast +from typing import TYPE_CHECKING, cast from uuid import UUID import pytest @@ -28,7 +28,9 @@ UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) -from pydantic import JsonValue + +if TYPE_CHECKING: + from pydantic import JsonValue def _plan() -> Stage12OutputPlan: @@ -94,7 +96,7 @@ def _spm_comparison() -> dict[str, JsonValue]: ), ) return cast( - dict[str, JsonValue], + "dict[str, JsonValue]", build_spm_comparison_provenance( baseline_receipts=(receipt,), reform_receipts=(receipt,), @@ -254,8 +256,8 @@ def test_aggregate_artifact_accepts_only_canonical_spm_provenance() -> None: ) mismatched = _spm_comparison() - reform = cast(dict[str, JsonValue], mismatched["reform"]) - receipt = cast(dict[str, JsonValue], reform["receipt"]) + reform = cast("dict[str, JsonValue]", mismatched["reform"]) + receipt = cast("dict[str, JsonValue]", reform["receipt"]) receipt["years"] = ["2027"] with pytest.raises(ValueError, match="baseline and reform"): AggregateReportArtifactPayload( diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/errors.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/errors.py index eb52a2750..e4bf391a3 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/errors.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/errors.py @@ -24,9 +24,10 @@ import logging import uuid -from typing import Any +from typing import TYPE_CHECKING, Any -from policyengine_observability import ObservabilityRuntime +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime logger = logging.getLogger(__name__) diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py index 7195be901..bdc96013e 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py @@ -3,9 +3,8 @@ import os from dataclasses import replace from importlib.metadata import PackageNotFoundError, version -from typing import Literal +from typing import TYPE_CHECKING, Literal -from fastapi import FastAPI from policyengine_observability import ( DeploymentIdentity, GoogleCloudLogDestination, @@ -21,6 +20,9 @@ process_instance_id, ) +if TYPE_CHECKING: + from fastapi import FastAPI + Platform = Literal["google_cloud_run", "modal", "local", "other"] MODAL_IMAGE_ENV_NAMES = ( diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py index b023ff0bb..c3101eaa7 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py @@ -7,10 +7,13 @@ from __future__ import annotations -from collections.abc import Mapping from dataclasses import dataclass from enum import StrEnum from types import MappingProxyType +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Mapping class RunConfiguration(StrEnum): 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 f463d0a32..fb6ca6e32 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py @@ -2,12 +2,9 @@ from __future__ import annotations -from collections.abc import Callable, Iterator from contextlib import contextmanager -from datetime import datetime from enum import StrEnum -from typing import Any, Protocol -from uuid import UUID +from typing import TYPE_CHECKING, Any, Protocol from policyengine_simulation_contract.stage12_execution import ( ComparisonReportPersistenceResult, @@ -43,6 +40,11 @@ require_successful_simulation_replay, ) +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + from datetime import datetime + from uuid import UUID + REPORT_IDENTITY_COLUMNS = ( "environment", "calculation_flow", diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py index dbaa5035c..2850f1ac4 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py @@ -3,6 +3,7 @@ from __future__ import annotations from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING from uuid import uuid4 from sqlalchemy import ( @@ -17,7 +18,6 @@ update, ) from sqlalchemy.dialects.postgresql import ENUM, UUID, insert -from sqlalchemy.engine import Connection, Engine from sqlalchemy.exc import SQLAlchemyError from policyengine_stage12_persistence.tables import ( @@ -31,6 +31,9 @@ comparison_simulations, ) +if TYPE_CHECKING: + from sqlalchemy.engine import Connection, Engine + REQUIRED_TABLE_PRIVILEGES = ("SELECT", "INSERT", "UPDATE", "DELETE") FORBIDDEN_TABLE_PRIVILEGES = ("TRUNCATE", "REFERENCES", "TRIGGER") diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validators.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validators.py index 0fe6e7ad1..6f6179e93 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validators.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validators.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Iterable +from typing import TYPE_CHECKING from policyengine_simulation_contract.stage12_execution import ( ComparisonReportRecord, @@ -12,6 +12,9 @@ ResultComparisonStatus, ) +if TYPE_CHECKING: + from collections.abc import Iterable + REPORT_IMMUTABLE_FIELDS = ( "evaluation_id", "contract_version", diff --git a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py index d17222c1a..726d5a741 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py @@ -323,7 +323,7 @@ def record_sleep(seconds: float) -> None: ) monkeypatch.setattr(time, "sleep", record_sleep) result = poll_for_completion( - cast(Client, object()), + cast("Client", object()), "fc-test", max_wait_seconds=max_wait_seconds, poll_interval=poll_interval, 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 287a0cb63..5e42810b1 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py @@ -5,16 +5,19 @@ import logging from collections.abc import Callable from functools import lru_cache -from typing import Literal +from typing import TYPE_CHECKING, Literal import jwt from fastapi import Depends, HTTPException, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer -from policyengine_observability import ObservabilityRuntime -from policyengine_simulation_entry.config import Settings from policyengine_simulation_entry.schemas import CallerIdentity +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + + from policyengine_simulation_entry.config import Settings + _bearer = HTTPBearer(auto_error=False) logger = logging.getLogger(__name__) 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 4a371b984..6cbdcbefd 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py @@ -5,7 +5,7 @@ import asyncio import time from dataclasses import dataclass -from typing import Protocol +from typing import TYPE_CHECKING, Protocol import httpx from policyengine_observability import ( @@ -13,13 +13,16 @@ ObservabilityRuntime, instrument_httpx, ) -from policyengine_simulation_contract.json_types import JsonObject from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER from pydantic import ValidationError -from policyengine_simulation_entry.config import Settings from policyengine_simulation_entry.schemas import BackendServiceToken +if TYPE_CHECKING: + from policyengine_simulation_contract.json_types import JsonObject + + from policyengine_simulation_entry.config import Settings + SAFE_RESPONSE_HEADERS = frozenset( { "cache-control", 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 c51e7aa2e..8095728cc 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 @@ -4,10 +4,9 @@ from dataclasses import dataclass from enum import StrEnum -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from uuid import UUID, uuid5 -from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ( BundleProvenance, DatasetArtifactMediaType, @@ -20,7 +19,10 @@ SimulationExecutionInput, SimulationRole, ) -from policyengine_simulation_contract.stage12_manifest import V2WorkerVersion + +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_bundle import CountryId + from policyengine_simulation_contract.stage12_manifest import V2WorkerVersion class ComparisonSkipReason(StrEnum): @@ -172,7 +174,7 @@ def adapt_annual_comparison( ) ) geography = GeographySelection( - country=cast(CountryId, country), + country=cast("CountryId", country), region=region, ) options = {key: payload[key] for key in ("spm",) if payload.get(key) is not None} diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py index f909e8455..a7afe6d03 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py @@ -6,11 +6,10 @@ import json import logging from datetime import UTC, datetime, timedelta -from typing import Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 import modal -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( ComparisonReportRecord, ComparisonRunAggregationStatus, @@ -34,9 +33,13 @@ ) from policyengine_stage12_persistence import Stage12PersistenceStore -from policyengine_simulation_entry.config import Settings from policyengine_simulation_entry.stage12_adapter import adapt_annual_comparison +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + + from policyengine_simulation_entry.config import Settings + logger = logging.getLogger(__name__) AUTOMATIC_REPORT_NAMESPACE = uuid5( NAMESPACE_URL, diff --git a/projects/policyengine-simulation-entry/tests/conftest.py b/projects/policyengine-simulation-entry/tests/conftest.py index 560c1d6ef..a250bacb6 100644 --- a/projects/policyengine-simulation-entry/tests/conftest.py +++ b/projects/policyengine-simulation-entry/tests/conftest.py @@ -1,14 +1,17 @@ from __future__ import annotations from dataclasses import dataclass +from typing import TYPE_CHECKING import pytest from fastapi.testclient import TestClient -from policyengine_simulation_contract.json_types import JsonObject from policyengine_simulation_entry.app import create_app from policyengine_simulation_entry.backend import BackendResponse from policyengine_simulation_entry.config import Settings +if TYPE_CHECKING: + from policyengine_simulation_contract.json_types import JsonObject + def make_settings(**overrides) -> Settings: values = { diff --git a/projects/policyengine-simulation-entry/tests/test_app.py b/projects/policyengine-simulation-entry/tests/test_app.py index d50e4e7f9..eb3835f19 100644 --- a/projects/policyengine-simulation-entry/tests/test_app.py +++ b/projects/policyengine-simulation-entry/tests/test_app.py @@ -3,6 +3,7 @@ import asyncio import json import logging +from typing import TYPE_CHECKING from unittest.mock import Mock from uuid import UUID @@ -10,7 +11,6 @@ from conftest import FakeBackend, make_settings from fastapi import HTTPException from policyengine_observability import REQUEST_ID_HEADER -from policyengine_simulation_contract.json_types import JsonObject from policyengine_simulation_contract.stage12_execution import ( ComparisonRunLifecycleStatus, SimulationRole, @@ -34,6 +34,9 @@ eligible_payload, ) +if TYPE_CHECKING: + from policyengine_simulation_contract.json_types import JsonObject + def response( status: int, diff --git a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py index 302133ae2..d158ff522 100644 --- a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py +++ b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py @@ -9,7 +9,7 @@ import json import re -from pathlib import Path +from typing import TYPE_CHECKING import pytest from workflow_script_support import ( @@ -18,6 +18,9 @@ run_script, ) +if TYPE_CHECKING: + from pathlib import Path + REUSABLE_DEPLOY_WORKFLOW = WORKFLOWS_DIR / "simulation-deploy.reusable.yml" diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 9c0ae2556..03146eb4e 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -105,7 +105,7 @@ def test_canonical_spm_extensions_are_public(): def test_route_table_is_frozen(): - assert methods(cast(OpenAPIDocument, create_app().openapi())) == EXPECTED_ROUTES + assert methods(cast("OpenAPIDocument", create_app().openapi())) == EXPECTED_ROUTES def test_old_gateway_routes_are_all_present(): @@ -117,10 +117,10 @@ def test_old_gateway_routes_are_all_present(): / "openapi.json" ) gateway_spec = cast( - OpenAPIDocument, + "OpenAPIDocument", json.loads(gateway_spec_path.read_text()), ) - cloud_run_routes = methods(cast(OpenAPIDocument, create_app().openapi())) + cloud_run_routes = methods(cast("OpenAPIDocument", create_app().openapi())) assert methods(gateway_spec) <= cloud_run_routes @@ -134,10 +134,10 @@ def test_normalized_contract_matches_old_gateway(): / "openapi.json" ) gateway_spec = cast( - OpenAPIDocument, + "OpenAPIDocument", json.loads(gateway_spec_path.read_text()), ) - cloud_run_spec = cast(OpenAPIDocument, create_app().openapi()) + cloud_run_spec = cast("OpenAPIDocument", create_app().openapi()) assert normalized_compatibility_paths( cloud_run_spec diff --git a/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py index 51257a71f..00b0b33be 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py @@ -4,12 +4,15 @@ import os import subprocess -from pathlib import Path +from typing import TYPE_CHECKING import pytest from fixtures.test_modal_scripts import REPO_ROOT, SCRIPTS_DIR +if TYPE_CHECKING: + from pathlib import Path + SCRIPT = SCRIPTS_DIR / "update-policyengine-package.sh" diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py index 5403375b0..e83286c69 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py @@ -2,15 +2,17 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_observability.stages import BUDGET_WINDOW_STAGES, Stage from src.modal.budget_window_context import build_batch_context from src.modal.budget_window_scheduler import BudgetWindowBatchRunner import modal +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + def run_budget_window_batch_impl( params: dict[str, Any], 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 602083f1c..2fe139224 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py @@ -3,9 +3,8 @@ from __future__ import annotations import time -from typing import Any +from typing import TYPE_CHECKING, Any -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.budget_window_state import ( build_batch_status_response, create_initial_batch_state, @@ -34,6 +33,9 @@ import modal +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + # 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 # Modal control-plane from getting hammered. We start aggressive (0.5s) so diff --git a/projects/policyengine-simulation-executor/src/modal/bundle_data.py b/projects/policyengine-simulation-executor/src/modal/bundle_data.py index f395a3235..97c78104f 100644 --- a/projects/policyengine-simulation-executor/src/modal/bundle_data.py +++ b/projects/policyengine-simulation-executor/src/modal/bundle_data.py @@ -3,7 +3,10 @@ from __future__ import annotations import shlex -from collections.abc import Sequence +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Sequence def bundle_data_install_command( diff --git a/projects/policyengine-simulation-executor/src/modal/fanout.py b/projects/policyengine-simulation-executor/src/modal/fanout.py index c138d1a3b..13ce29bcb 100644 --- a/projects/policyengine-simulation-executor/src/modal/fanout.py +++ b/projects/policyengine-simulation-executor/src/modal/fanout.py @@ -10,8 +10,10 @@ from __future__ import annotations -from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Iterable def build_child_payload( diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index b769a6181..1aa1a8068 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -27,9 +27,8 @@ import time from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.job_progress import JobExecution, JobStage from policyengine_simulation_contract.spm import ( SPMInputError, @@ -56,6 +55,9 @@ import modal +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + logger = logging.getLogger(__name__) POLL_INTERVAL_INITIAL_SECONDS = 0.5 diff --git a/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py b/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py index f7715d5e5..a5d9fad46 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py @@ -10,12 +10,15 @@ import argparse import os -from collections.abc import Mapping from datetime import datetime, timezone +from typing import TYPE_CHECKING from policyengine_simulation_executor.artifact_store import ArtifactStore from policyengine_simulation_executor.precompute_models import DeployedMarker +if TYPE_CHECKING: + from collections.abc import Mapping + def build_marker( environment: str, diff --git a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py index fc088bd2e..fb96e8c07 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py @@ -6,11 +6,10 @@ import json from datetime import datetime, timezone from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import uuid4 import modal -from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_manifest import ( ManifestText, V2WorkerValidation, @@ -18,6 +17,9 @@ ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_bundle import CountryId + def _spawn_validation( *, diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 076f9dce0..f79c96578 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -10,8 +10,8 @@ import json import os from pathlib import Path +from typing import TYPE_CHECKING -from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_manifest import v2_application_name from policyengine_simulation_executor.hf_access_validation import ( HF_MODAL_SECRET_NAME, @@ -40,6 +40,9 @@ import modal +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_bundle import CountryId + STAGE12_DATA_DIR = "/opt/policyengine/stage12-data" _UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "." RESOLVED_BUNDLE = load_stage12_bundle() 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 e58707a46..f8e2caa78 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 collections.abc import Callable +from typing import TYPE_CHECKING from policyengine_simulation_executor.artifact_keys import ( BaselineArtifactIdentity, @@ -54,6 +54,9 @@ WorkSelection, ) +if TYPE_CHECKING: + from collections.abc import Callable + logger = logging.getLogger(__name__) PRECOMPUTE_COUNTRY = "us" 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 db9106831..bdeff1126 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 @@ -4,11 +4,9 @@ from contextlib import nullcontext from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.job_progress import JobStage -from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_observability.stages import ( ANNUAL_IMPACT_STAGES, Stage, @@ -45,6 +43,10 @@ SingleYearMacroOutput, ) +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata + @dataclass class SimulationOutputBuilder: 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 3e71d8df5..4deb42a05 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 @@ -10,7 +10,6 @@ from policyengine.data.uk_geography_assets import ( CONSTITUENCY_ASSET_SPEC, ) -from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_executor.simulation_macro_output import ( CongressionalDistrictImpactOutput, @@ -26,6 +25,7 @@ if TYPE_CHECKING: from policyengine.data.uk_geography_assets import UKGeographyAssetSpec + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata def _required_uk_geography_lookup_csv_path(spec: UKGeographyAssetSpec) -> str: 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 3e70637da..de63db21d 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 @@ -2,8 +2,7 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping -from typing import Any +from typing import TYPE_CHECKING, Any from policyengine_simulation_executor.simulation_macro_output import ( AgePovertyOutput, @@ -22,6 +21,9 @@ _try_compute_output, ) +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping + US_POVERTY_TYPES = { "spm": "poverty", "spm_deep": "deep_poverty", 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 5b07ffa1c..798c091af 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 @@ -3,12 +3,9 @@ from __future__ import annotations import json -from collections.abc import Mapping -from datetime import datetime from hashlib import sha256 from io import BytesIO -from typing import Any -from uuid import UUID +from typing import TYPE_CHECKING, Any import pandas as pd import pyarrow as pa @@ -29,6 +26,11 @@ from policyengine_simulation_executor.artifact_store import ArtifactStore +if TYPE_CHECKING: + from collections.abc import Mapping + from datetime import datetime + from uuid import UUID + SIMULATION_MEDIA_TYPE = ArtifactMediaType.PARQUET.value REPORT_MEDIA_TYPE = ArtifactMediaType.JSON.value PARQUET_CONTRACT = SIMULATION_PARQUET_PAYLOAD_CONTRACT diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/assertions.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/assertions.py index e515c88c5..01d090442 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/assertions.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/assertions.py @@ -2,12 +2,15 @@ from __future__ import annotations -from collections.abc import Mapping - -from policyengine_simulation_contract.stage12_bundle import Stage12BundleManifest +from typing import TYPE_CHECKING from .errors import Stage12BundleError +if TYPE_CHECKING: + from collections.abc import Mapping + + from policyengine_simulation_contract.stage12_bundle import Stage12BundleManifest + def assertion_values(bundle: Stage12BundleManifest) -> dict[str, str]: values = { diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py index 93c3bd737..87adc695b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py @@ -3,8 +3,8 @@ from __future__ import annotations import json -from collections.abc import Sequence from hashlib import sha256 +from typing import TYPE_CHECKING from policyengine_simulation_contract.stage12_bundle import ( SUPPORTED_COUNTRIES, @@ -22,6 +22,9 @@ require_text, ) +if TYPE_CHECKING: + from collections.abc import Sequence + def _normalize_stage12_bundle( raw_bundle: object, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py index 335943bb0..29d18a676 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py @@ -5,8 +5,7 @@ from datetime import UTC, datetime from decimal import Decimal from hashlib import sha256 -from typing import TypeGuard -from uuid import UUID +from typing import TYPE_CHECKING, TypeGuard from policyengine_simulation_contract.stage12_execution import ( ResultComparisonArtifactPayload, @@ -16,6 +15,9 @@ from policyengine_simulation_executor.stage12_artifacts import canonical_json_bytes +if TYPE_CHECKING: + from uuid import UUID + type ResultObject = dict[str, JsonValue] _result_adapter = TypeAdapter(ResultObject) NUMERIC_ABSOLUTE_TOLERANCE = Decimal("0.01") 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 a9dadb23a..b61a3c0f1 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 @@ -2,19 +2,22 @@ from __future__ import annotations -from collections.abc import Mapping -from typing import Any +from typing import TYPE_CHECKING, Any -import pandas as pd -from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ( ReportExecutionInput, SimulationArtifactDescriptor, Stage12OutputPlan, stage12_output_plan_sha256, ) -from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata -from pydantic import JsonValue + +if TYPE_CHECKING: + from collections.abc import Mapping + + import pandas as pd + from policyengine_simulation_contract.stage12_bundle import CountryId + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata + from pydantic import JsonValue def validate_uk_local_authority_metadata( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py index c61bdd50c..93c35e3d7 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py @@ -5,9 +5,8 @@ import logging from contextlib import nullcontext from datetime import UTC, datetime -from typing import Any +from typing import TYPE_CHECKING, Any -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( ComparisonReportRecord, ResultComparisonStatus, @@ -15,12 +14,17 @@ ) from policyengine_simulation_observability.stages import Stage, StagePlan -from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore from policyengine_simulation_executor.stage12_result_comparison import compare_results -from .dependencies import ChildInvoker, ComparisonStore from .failure_causes import failure_summary +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + + from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore + + from .dependencies import ChildInvoker, ComparisonStore + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS = 900 logger = logging.getLogger(__name__) 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 ae1824909..e6dfc569c 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 @@ -2,14 +2,12 @@ from __future__ import annotations -from collections.abc import Callable from datetime import UTC, datetime from functools import partial from hashlib import sha256 -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import UUID, uuid4 -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( AggregateReportArtifactDescriptor, ComparisonReportRecord, @@ -65,6 +63,11 @@ ) from .simulation import descriptor_from_record, simulation_input_sha256 +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_observability import ObservabilityRuntime + SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py index e5848ebe0..f3d87885f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py @@ -3,24 +3,27 @@ from __future__ import annotations import os -from datetime import datetime from functools import partial -from typing import Any, Protocol -from uuid import UUID - -from policyengine_observability import ObservabilityRuntime -from policyengine_simulation_contract.stage12_execution import ( - ComparisonReportPersistenceResult, - ComparisonReportRecord, - ComparisonSimulationPersistenceResult, - ComparisonSimulationRecord, -) +from typing import TYPE_CHECKING, Any, Protocol + from policyengine_stage12_persistence import Stage12PersistenceStore from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore from .database_retry import open_connection_with_retry +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.stage12_execution import ( + ComparisonReportPersistenceResult, + ComparisonReportRecord, + ComparisonSimulationPersistenceResult, + ComparisonSimulationRecord, + ) + class ComparisonStore(Protocol): def create_or_resolve_report( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py index 4d899d6c6..ed2731dca 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py @@ -4,16 +4,14 @@ from collections.abc import Callable, Mapping from contextlib import AbstractContextManager, nullcontext -from typing import Protocol, cast +from typing import TYPE_CHECKING, Protocol, cast -import pandas as pd from policyengine.core import Simulation from policyengine.outputs import ( configure_cliff_impact_variables, configure_labor_supply_response_variables, labor_supply_response_is_active, ) -from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ( EntityOutputPlan, PlannedSimulationExecutionInput, @@ -25,6 +23,10 @@ ) from policyengine_simulation_observability.stages import Stage +if TYPE_CHECKING: + import pandas as pd + from policyengine_simulation_contract.stage12_bundle import CountryId + UK_GEOGRAPHIC_DATASET_VARIABLES: dict[str, tuple[str, ...]] = { "household": ("constituency_code_oa", "la_code_oa"), } @@ -230,7 +232,7 @@ def apply_output_plan( if variable not in entity_extras: entity_extras.append(variable) simulation.extra_variables = extras - model = cast(OutputVariableModel, simulation.tax_benefit_model_version) + model = cast("OutputVariableModel", simulation.tax_benefit_model_version) resolved = model.resolve_entity_variables(simulation) missing = { entity_plan.entity: sorted( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index cd3c9cc30..025b22c9d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -7,12 +7,9 @@ from datetime import UTC, datetime from functools import partial from hashlib import sha256 -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd -from policyengine_observability import ObservabilityRuntime -from policyengine_simulation_contract.spm import SPMProvenance -from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ( ComparisonRunLifecycleStatus, ComparisonSimulationRecord, @@ -22,7 +19,6 @@ Stage12InvocationContext, stage12_output_plan_sha256, ) -from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_observability.stages import ( STAGE12_SIMULATION_STAGES, Stage, @@ -47,6 +43,12 @@ from .failure_recording import persist_failure_record, report_original_failure from .output_planning import apply_output_plan, validate_output_frames +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.spm import SPMProvenance + from policyengine_simulation_contract.stage12_bundle import CountryId + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata + @dataclass(frozen=True) class SimulationCalculation: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index 008bfae40..0175003cb 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -7,15 +7,16 @@ from hashlib import sha256 from importlib import import_module from pathlib import Path -from typing import Any - -from policyengine_simulation_contract.stage12_bundle import CountryId +from typing import TYPE_CHECKING, Any from policyengine_simulation_executor.release_bundle import ( resolve_local_bundle_dataset_path, ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_bundle import CountryId + REQUIRED_SECRET_ALTERNATIVES = ( ("HUGGING_FACE_TOKEN",), ( diff --git a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py index 9ce66bebe..ba5f0522f 100644 --- a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py @@ -4,7 +4,7 @@ import subprocess import sys -from pathlib import Path +from typing import TYPE_CHECKING import pytest from fixtures.test_policyengine_package_update_scripts import ( @@ -16,6 +16,9 @@ updater_env, ) +if TYPE_CHECKING: + from pathlib import Path + pytest_plugins = ("fixtures.test_policyengine_package_update_scripts",) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py index 4c9b92301..f22cb7e8f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py @@ -3,12 +3,12 @@ from __future__ import annotations from contextlib import contextmanager +from typing import TYPE_CHECKING from uuid import UUID import pandas as pd import pytest from policyengine.core import Simulation -from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ( BundleProvenance, DatasetArtifactMediaType, @@ -28,7 +28,10 @@ validate_output_frames, ) from policyengine_simulation_observability.stages import Stage -from pydantic import JsonValue + +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_bundle import CountryId + from pydantic import JsonValue EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") diff --git a/projects/policyengine-simulation-gateway/fixtures/package_imports.py b/projects/policyengine-simulation-gateway/fixtures/package_imports.py index 8c93aafa2..b6769a469 100644 --- a/projects/policyengine-simulation-gateway/fixtures/package_imports.py +++ b/projects/policyengine-simulation-gateway/fixtures/package_imports.py @@ -4,11 +4,14 @@ import importlib import sys -from collections.abc import Iterator from dataclasses import dataclass +from typing import TYPE_CHECKING import pytest +if TYPE_CHECKING: + from collections.abc import Iterator + GATEWAY_MODEL_MODULE = "policyengine_simulation_contract.gateway_models" GATEWAY_ENDPOINTS_MODULE = "policyengine_simulation_gateway.endpoints" GATEWAY_PACKAGE_MODULE = "policyengine_simulation_gateway" 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 db926a45a..620e7f774 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py @@ -2,12 +2,16 @@ from __future__ import annotations -from fastapi import FastAPI, Request +from typing import TYPE_CHECKING + from policyengine_simulation_observability.identifiers import ( OBSERVABILITY_ID_HEADER, normalize_observability_id, ) +if TYPE_CHECKING: + from fastapi import FastAPI, Request + def install_observability_id_middleware(app: FastAPI) -> None: """Store an incoming identifier candidate without creating a workflow.""" 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 e9c2c5f3d..c658c0e32 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py @@ -2,10 +2,14 @@ from __future__ import annotations +from typing import TYPE_CHECKING + from fastapi.responses import JSONResponse -from policyengine_simulation_contract.gateway_models import ( - BudgetWindowBatchStatusResponse, -) + +if TYPE_CHECKING: + from policyengine_simulation_contract.gateway_models import ( + BudgetWindowBatchStatusResponse, + ) class ServerErrorResponse(JSONResponse): diff --git a/ruff.toml b/ruff.toml index 482af947b..a22dc8cc2 100644 --- a/ruff.toml +++ b/ruff.toml @@ -11,4 +11,40 @@ select = [ "E4", "E7", "E9", "F", # Group 1: safe auto-fixes. "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W", + # Group 2: typing-only imports. + "TC", +] + +[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 +# class that other models inherit from is listed. +runtime-evaluated-base-classes = [ + "policyengine_fastapi.observability.contracts.CorrelatedRunFields", + "policyengine_fastapi.observability.contracts.ObservabilityModel", + "policyengine_simulation_contract.gateway_models.GatewayRequestBase", + "policyengine_simulation_contract.macro_output.MacroOutputModel", + "policyengine_simulation_contract.macro_output.MacroRootModel", + "policyengine_simulation_contract.stage12_bundle.StrictBundleModel", + "policyengine_simulation_contract.stage12_execution.SimulationExecutionInput", + "policyengine_simulation_contract.stage12_execution.StrictContractModel", + "policyengine_simulation_contract.stage12_manifest.StrictManifestModel", + "policyengine_simulation_executor.precompute_models._StrictModel", + "pydantic.BaseModel", + "pydantic.RootModel", +] +# FastAPI evaluates endpoint annotations when routes are registered. +runtime-evaluated-decorators = [ + "fastapi.APIRouter.api_route", + "fastapi.APIRouter.delete", + "fastapi.APIRouter.get", + "fastapi.APIRouter.patch", + "fastapi.APIRouter.post", + "fastapi.APIRouter.put", + "fastapi.FastAPI.api_route", + "fastapi.FastAPI.delete", + "fastapi.FastAPI.get", + "fastapi.FastAPI.patch", + "fastapi.FastAPI.post", + "fastapi.FastAPI.put", ] diff --git a/scripts/rules/baseline.py b/scripts/rules/baseline.py index 9728f904b..f1a8aeb4b 100644 --- a/scripts/rules/baseline.py +++ b/scripts/rules/baseline.py @@ -4,11 +4,14 @@ import os import subprocess -from pathlib import Path +from typing import TYPE_CHECKING from scripts.rules.load import REGISTRY_PATH, load_text from scripts.rules.model import Registry, Status, Violation +if TYPE_CHECKING: + from pathlib import Path + class BaseUnavailable(Exception): """The base ref could not be read.""" diff --git a/scripts/rules/discover.py b/scripts/rules/discover.py index c7c7e42d9..323425618 100644 --- a/scripts/rules/discover.py +++ b/scripts/rules/discover.py @@ -8,7 +8,7 @@ from __future__ import annotations -from pathlib import Path +from typing import TYPE_CHECKING from scripts.rules.model import Kind, Registry, Violation from scripts.rules.resolve import ( @@ -22,6 +22,9 @@ workflow_jobs, ) +if TYPE_CHECKING: + from pathlib import Path + def _hidden(message: str, fix: str) -> Violation: return Violation(None, "hidden-rule", message, fix) diff --git a/scripts/rules/model.py b/scripts/rules/model.py index 1e7c68739..619138b01 100644 --- a/scripts/rules/model.py +++ b/scripts/rules/model.py @@ -2,11 +2,13 @@ from __future__ import annotations -from collections.abc import Mapping from dataclasses import dataclass, field -from datetime import date from enum import StrEnum -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Mapping + from datetime import date class Tier(StrEnum): diff --git a/scripts/rules/resolve.py b/scripts/rules/resolve.py index 1e47b09c7..cddb6536d 100644 --- a/scripts/rules/resolve.py +++ b/scripts/rules/resolve.py @@ -10,14 +10,16 @@ import ast import re import tomllib -from collections.abc import Callable from functools import cache -from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from scripts.rules.model import CiRef, Kind, Rule, Violation from scripts.rules.workflows import JobBlock, parse_jobs +if TYPE_CHECKING: + from collections.abc import Callable + from pathlib import Path + DEFAULT_RUNNER = "scripts/run_quality_guards.py" RUFF_DEFAULT_SELECT = ("E4", "E7", "E9", "F") PROOF_CI_FIELDS = ("workflow", "job", "invocation") diff --git a/scripts/rules/validate.py b/scripts/rules/validate.py index 0469a5260..f7d117d65 100644 --- a/scripts/rules/validate.py +++ b/scripts/rules/validate.py @@ -4,12 +4,14 @@ import re from datetime import date -from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from scripts.rules.model import Kind, Registry, Rule, Status, Tier, Violation from scripts.rules.resolve import node_exists +if TYPE_CHECKING: + from pathlib import Path + RULE_ID = re.compile(r"^[a-z0-9]+(\.[a-z0-9][a-z0-9-]*)+$") # Required and optional enforcer fields per kind, with their expected types. From 57bcda8557d39476ed65a1b5cdf0ac53d10c8c4b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:39:18 +0400 Subject: [PATCH 02/10] Register TC in the enforced ruff lint rule Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- changelog_entry.yaml | 1 + docs/generated/rules.md | 8 ++++---- rules/registry.toml | 2 +- 4 files changed, 7 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 26122fae6..62885a850 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -139,7 +139,7 @@ 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, 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 (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`) ### CI workflows and scripts diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 83047a7a5..e1109ee7e 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 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 - Moved multi-line `run` blocks in the PR, image smoke and client publishing workflows into single-line steps and tested scripts, tested the dependency-update scripts, fixed the shellcheck finding in modal-precompute.sh, and enforced the CI-script testing and shellcheck rules with empty baselines diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 620a8e7a6..ed76deec6 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -249,7 +249,7 @@ 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, 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 (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` | — | ```mermaid @@ -261,7 +261,7 @@ 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__UP__W_["ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, UP, W)"] + 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)"] 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"] @@ -269,8 +269,8 @@ flowchart LR 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__UP__W_ - e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__UP__W_ --> j_CI_pr_yml_lint + 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_tests_pass --> e_CI_job_test e_CI_job_test --> j_CI_pr_yml_test ``` diff --git a/rules/registry.toml b/rules/registry.toml index 605f0fa22..93075414e 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -732,7 +732,7 @@ autofix = true [rule.enforcer] kind = "ruff" mode = "lint" -codes = ["E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W"] +codes = ["E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"] [[rule.ci]] workflow = ".github/workflows/pr.yml" job = "lint" From 3e808b024b190420c1f837ec762f9861b8f70de3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:16:26 +0400 Subject: [PATCH 03/10] Move the typing-only imports #789 added under TYPE_CHECKING #789 landed on main after this stack branched. TC, which this PR enables, flags the annotation-only imports in its new database-retry and failure-recording modules; they now load only under TYPE_CHECKING. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/database_retry.py | 10 +++++++--- .../stage12_runtime/failure_recording.py | 7 +++++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py index a432cf3fe..cbaed0477 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py @@ -5,13 +5,17 @@ import logging import random import time -from collections.abc import Callable from contextlib import suppress +from typing import TYPE_CHECKING -from policyengine_observability import ObservabilityRuntime -from sqlalchemy.engine import Connection from sqlalchemy.exc import OperationalError +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_observability import ObservabilityRuntime + from sqlalchemy.engine import Connection + logger = logging.getLogger(__name__) # Coordinators and workers start in bursts against a session pooler with a diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py index b86c0679c..7688aef5c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py @@ -3,10 +3,13 @@ from __future__ import annotations import logging -from collections.abc import Callable, Mapping from contextlib import suppress +from typing import TYPE_CHECKING -from policyengine_observability import ObservabilityRuntime +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + + from policyengine_observability import ObservabilityRuntime logger = logging.getLogger(__name__) From 2f2c8748d476aa9fc4bff423d76813b1c940429b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 02:09:19 +0400 Subject: [PATCH 04/10] Move the typing-only Callable import #795 added under TYPE_CHECKING Co-Authored-By: Claude Opus 5.5 --- .../src/modal/segmented_national.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 1aa1a8068..94b3d5ef3 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -25,7 +25,6 @@ import logging import time -from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -56,6 +55,8 @@ import modal if TYPE_CHECKING: + from collections.abc import Callable + from policyengine_observability import ObservabilityRuntime logger = logging.getLogger(__name__) From 51c5f414f15b42cfaebe77cf1dc4c0303698d51f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 02:34:34 +0400 Subject: [PATCH 05/10] Move the typing-only imports #791 added under TYPE_CHECKING #791 landed on main after this stack branched. TC, which this PR enables, flags the imports its new Stage 12 failure-cause module uses only in annotations (Iterator, Mapping, ObservabilityRuntime, Stage and StagePlan); they move under TYPE_CHECKING. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/failure_causes.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py index b6e20b191..e96686a6e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py @@ -29,18 +29,22 @@ import logging import os import re -from collections.abc import Iterator, Mapping from contextlib import AbstractContextManager, contextmanager, nullcontext +from typing import TYPE_CHECKING from urllib.parse import urlsplit -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ErrorSummary -from policyengine_simulation_observability.stages import Stage, StagePlan from pydantic import TypeAdapter, ValidationError from sqlalchemy.exc import StatementError from policyengine_simulation_executor.hf_access_validation import HF_RUNTIME_ENV_NAME +if TYPE_CHECKING: + from collections.abc import Iterator, Mapping + + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_observability.stages import Stage, StagePlan + # The longest ``error_summary`` the contract accepts. MAX_FAILURE_SUMMARY_LENGTH: int = TypeAdapter(ErrorSummary).json_schema()["maxLength"] # Redaction reads at most this much of a message. The cap bounds the cost of From 2c9f5fbbb64aea3708cb34685ff59144d0197c98 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:51:58 +0400 Subject: [PATCH 06/10] List the Pydantic bases ruff reads through imports as runtime-evaluated Ruff treats a class as runtime-evaluated when a base, read through the import statements of the class's own module, is listed in runtime-evaluated-base-classes, or is a class in that same module that is. It never follows a re-export, and only follows inheritance within a module. Four models had bases ruff reads at unlisted paths, so ruff would have moved imports their annotations need under TYPE_CHECKING if they gained any: - ArtifactBaselineSimulation and PrecomputedSimulation subclass policyengine.py's policyengine.core.Simulation; - AppSettings subclasses pydantic_settings.BaseSettings; - EconomyComparison subclassed SingleYearMacroOutput through the executor's simulation_macro_output re-export. The first two paths are listed as their subclasses import them. For EconomyComparison, compat_models now imports SingleYearMacroOutput from policyengine_simulation_contract.macro_output, the module that defines it, and that path is listed: every project entry then names its defining module, and the entry also covers subclasses elsewhere that import the class from the contract. The ruff.toml comment now describes how ruff resolves bases instead of "direct bases only". Co-Authored-By: Claude Opus 5.5 --- .../policyengine_simulation_executor/compat_models.py | 5 +---- ruff.toml | 11 +++++++++-- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/compat_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/compat_models.py index beaca988d..0c8ec8016 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/compat_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/compat_models.py @@ -4,13 +4,10 @@ from typing import Any +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput from policyengine_simulation_contract.spm import SPMSelection from pydantic import BaseModel, ConfigDict -from policyengine_simulation_executor.simulation_macro_output import ( - SingleYearMacroOutput, -) - class SimulationOptions(BaseModel): """Legacy request schema name kept for generated clients.""" diff --git a/ruff.toml b/ruff.toml index a22dc8cc2..88b366afc 100644 --- a/ruff.toml +++ b/ruff.toml @@ -17,14 +17,20 @@ select = [ [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 -# class that other models inherit from is listed. +# stay at runtime. Ruff reads each base through the imports of the subclass's +# own module and follows inheritance only through classes in that module, never +# through a re-export, so a base is listed by the path its subclasses import it +# from: the defining module for project bases, the public path for third-party +# ones. See docs/engineering/skills/typing-only-imports.md; the +# sim.quality.runtime-evaluated-models-resolve tests check every model. runtime-evaluated-base-classes = [ + "policyengine.core.Simulation", "policyengine_fastapi.observability.contracts.CorrelatedRunFields", "policyengine_fastapi.observability.contracts.ObservabilityModel", "policyengine_simulation_contract.gateway_models.GatewayRequestBase", "policyengine_simulation_contract.macro_output.MacroOutputModel", "policyengine_simulation_contract.macro_output.MacroRootModel", + "policyengine_simulation_contract.macro_output.SingleYearMacroOutput", "policyengine_simulation_contract.stage12_bundle.StrictBundleModel", "policyengine_simulation_contract.stage12_execution.SimulationExecutionInput", "policyengine_simulation_contract.stage12_execution.StrictContractModel", @@ -32,6 +38,7 @@ runtime-evaluated-base-classes = [ "policyengine_simulation_executor.precompute_models._StrictModel", "pydantic.BaseModel", "pydantic.RootModel", + "pydantic_settings.BaseSettings", ] # FastAPI evaluates endpoint annotations when routes are registered. runtime-evaluated-decorators = [ From c6951ed248d262e44aebe01948d51e56923c5b74 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:51:58 +0400 Subject: [PATCH 07/10] Check that ruff treats every Pydantic model as runtime-evaluated The executor's test imports every module of the executor package and of the project libraries installed with it; the entrypoint's imports every module of its package. For every Pydantic model a module-level class statement defines, they check that: - ruff treats it as runtime-evaluated, resolving each base the way ruff does: through the model module's import statements as they stand at the class statement, reading Base[Item] as Base, following bases defined in the same module, and never following a re-export or an assignment. A failure names the model and the base path to list; - its annotations resolve, with model_rebuild(raise_errors=True). Unit tests pin the resolution: a listed base, a subscripted one, a relative import, a same-module chain through Generic, an unlisted re-export, an unlisted third-party base, an assigned alias, and a base in a package's __init__.py, which ruff names with __init__ in its path. The proof fixtures are compat_models.py before and after the previous commit. The executor's src/modal Modal entry files are not imported: they are not part of the installed package, and importing them builds Modal images and functions, which their tests do only with a fake Modal SDK and a patched environment. A static test checks that every class there derives only from builtins and the standard library, so none can be a model. The gateway defines no Pydantic models. The entrypoint shares no package with the executor that could hold test helpers, so it keeps a copy of the resolution helpers. Co-Authored-By: Claude Opus 5.5 --- .../test_runtime_evaluated_annotations.py | 399 +++++++++++++++ .../test_runtime_evaluated_annotations.py | 463 ++++++++++++++++++ .../bad_compat_models.py | 16 + .../good_compat_models.py | 13 + 4 files changed, 891 insertions(+) create mode 100644 projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py create mode 100644 projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py create mode 100644 rules/fixtures/sim.quality.runtime-evaluated-models-resolve/bad_compat_models.py create mode 100644 rules/fixtures/sim.quality.runtime-evaluated-models-resolve/good_compat_models.py diff --git a/projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py b/projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py new file mode 100644 index 000000000..51c17d1f6 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py @@ -0,0 +1,399 @@ +"""Pydantic models keep the runtime imports their annotations need. + +The packages' modules use ``from __future__ import annotations``, so ruff's +typing-only import rules (TC) move an import that only annotations use under +``if TYPE_CHECKING:``. Pydantic evaluates model annotations at runtime, so ruff +must treat every model as runtime-evaluated. It does when a base of the model, +read through the imports of the model's own module, is listed in +``runtime-evaluated-base-classes`` in ruff.toml, or is a class defined in that +same module that is runtime-evaluated itself. Ruff never follows a re-export to +the module that defines a class, so these checks resolve bases exactly as ruff +does. See docs/engineering/skills/typing-only-imports.md. + +The resolution helpers are a copy of those in the executor's +tests/test_runtime_evaluated_annotations.py, because the two projects share no +package that test helpers could live in; keep the copies the same. The +executor's copy checks the project libraries. +""" + +from __future__ import annotations + +import ast +import builtins +import dataclasses +import importlib +import importlib.util +import pkgutil +from functools import cache +from pathlib import Path +from typing import TYPE_CHECKING + +import pydantic +import pytest +import tomllib + +if TYPE_CHECKING: + from collections.abc import Iterator + from types import ModuleType + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +CHECKED_PACKAGES = ("policyengine_simulation_entry",) + +# What a module-level name is bound to: the qualified name an import gives it, +# a class statement in the same module, or None when ruff cannot resolve it (an +# assignment or a function definition). +Binding = str | ast.ClassDef | None + + +@dataclasses.dataclass(frozen=True) +class ModuleClass: + """A module-level class statement, with the names bound where it runs.""" + + module_name: str + ruff_module_path: str + node: ast.ClassDef + bindings: dict[str, Binding] + + @property + def qualified_name(self) -> str: + return f"{self.module_name}.{self.node.name}" + + +@dataclasses.dataclass(frozen=True) +class ResolvedBase: + """A base as ruff reads it, and whether ruff follows it within the module.""" + + written: str + path: str | None + same_module: bool + + +def _module_level_statements(body: list[ast.stmt]) -> Iterator[ast.stmt]: + """Yield the statements that run at module level, in source order.""" + for statement in body: + yield statement + if isinstance(statement, ast.If): + yield from _module_level_statements(statement.body) + yield from _module_level_statements(statement.orelse) + elif isinstance(statement, (ast.Try, ast.TryStar)): + for block in ( + statement.body, + *(handler.body for handler in statement.handlers), + statement.orelse, + statement.finalbody, + ): + yield from _module_level_statements(block) + + +def _names_bound_by(statement: ast.stmt, package: str) -> dict[str, Binding]: + """Return the module-level names a statement binds, as ruff resolves them.""" + if isinstance(statement, ast.Import): + bound: dict[str, Binding] = {} + for alias in statement.names: + if alias.asname is None: + top_level = alias.name.partition(".")[0] + bound[top_level] = top_level + else: + bound[alias.asname] = alias.name + return bound + if isinstance(statement, ast.ImportFrom): + source_module = ( + importlib.util.resolve_name( + "." * statement.level + (statement.module or ""), package + ) + if statement.level + else statement.module + ) + return { + alias.asname or alias.name: f"{source_module}.{alias.name}" + for alias in statement.names + if alias.name != "*" + } + if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef)): + return {statement.name: None} + if isinstance(statement, ast.Assign) or ( + isinstance(statement, ast.AnnAssign) and statement.value is not None + ): + targets = ( + statement.targets + if isinstance(statement, ast.Assign) + else [statement.target] + ) + return { + node.id: None + for target in targets + for node in ast.walk(target) + if isinstance(node, ast.Name) + } + return {} + + +def module_classes( + module_name: str, source: str, *, is_package: bool +) -> dict[str, ModuleClass]: + """Return each module-level class with the names bound at its statement. + + Ruff names a module after its file, so a package's ``__init__.py`` is + ``package.__init__`` to ruff. + """ + package = module_name if is_package else module_name.rpartition(".")[0] + ruff_module_path = f"{module_name}.__init__" if is_package else module_name + bindings: dict[str, Binding] = {} + classes: dict[str, ModuleClass] = {} + for statement in _module_level_statements(ast.parse(source).body): + if isinstance(statement, ast.ClassDef): + classes[statement.name] = ModuleClass( + module_name, ruff_module_path, statement, dict(bindings) + ) + bindings[statement.name] = statement + else: + bindings.update(_names_bound_by(statement, package)) + return classes + + +def _dotted_name(base: ast.expr) -> list[str] | None: + """Split ``a.b.Base`` into its names, reading ``Base[Item]`` as ``Base``.""" + if isinstance(base, ast.Subscript): + base = base.value + attributes: list[str] = [] + while isinstance(base, ast.Attribute): + attributes.insert(0, base.attr) + base = base.value + if not isinstance(base, ast.Name): + return None + return [base.id, *attributes] + + +def _ruff_path(module_class: ModuleClass, base: ast.expr) -> str | None: + """Return the qualified name ruff resolves a base to, through the imports.""" + names = _dotted_name(base) + if names is None: + return None + head, *attributes = names + if head not in module_class.bindings: + return ".".join(["builtins", *names]) if hasattr(builtins, head) else None + binding = module_class.bindings[head] + if binding is None: + return None + if isinstance(binding, ast.ClassDef): + binding = f"{module_class.ruff_module_path}.{head}" + return ".".join([binding, *attributes]) + + +def resolved_bases( + module_class: ModuleClass, followed: frozenset[str] = frozenset() +) -> list[ResolvedBase]: + """Return the bases ruff checks, following classes in the same module. + + Ruff resolves a same-module base's own bases with the names bound at the + subclass, so the subclass's bindings are kept on the way down. + """ + found: list[ResolvedBase] = [] + for base in module_class.node.bases: + names = _dotted_name(base) + local_class = ( + module_class.bindings.get(names[0]) + if names is not None and len(names) == 1 + else None + ) + same_module = isinstance(local_class, ast.ClassDef) + found.append( + ResolvedBase(ast.unparse(base), _ruff_path(module_class, base), same_module) + ) + if isinstance(local_class, ast.ClassDef) and local_class.name not in followed: + found += resolved_bases( + dataclasses.replace(module_class, node=local_class), + followed | {local_class.name}, + ) + return found + + +def is_runtime_evaluated(module_class: ModuleClass, listed: frozenset[str]) -> bool: + """Whether ruff treats the class's annotations as evaluated at runtime.""" + return any(base.path in listed for base in resolved_bases(module_class)) + + +def _names_a_model(path: str) -> bool: + try: + target = pkgutil.resolve_name(path) + except (AttributeError, ImportError, ValueError): + return False + return isinstance(target, type) and issubclass(target, pydantic.BaseModel) + + +def _bases_to_list(module_class: ModuleClass) -> list[str]: + """Return the model bases from other modules that would fix the class.""" + outside = [ + base.path or base.written + for base in resolved_bases(module_class) + if not base.same_module + ] + return [path for path in outside if _names_a_model(path)] or outside + + +def _listed_runtime_evaluated_bases() -> frozenset[str]: + with (REPOSITORY_ROOT / "ruff.toml").open("rb") as ruff_config: + settings = tomllib.load(ruff_config) + type_checking = settings["lint"]["flake8-type-checking"] + return frozenset(type_checking["runtime-evaluated-base-classes"]) + + +@cache +def _imported_modules() -> tuple[ModuleType, ...]: + modules = [] + for package_name in CHECKED_PACKAGES: + package = importlib.import_module(package_name) + modules.append(package) + modules += [ + importlib.import_module(module_info.name) + for module_info in pkgutil.walk_packages( + package.__path__, f"{package_name}." + ) + ] + return tuple(modules) + + +@cache +def _model_classes() -> tuple[tuple[ModuleClass, type[pydantic.BaseModel]], ...]: + """Return each Pydantic model a module-level class statement defines.""" + models = [] + for module in _imported_modules(): + source_path = Path(module.__file__) + classes = module_classes( + module.__name__, + source_path.read_text(encoding="utf-8"), + is_package=source_path.name == "__init__.py", + ) + for class_name, module_class in classes.items(): + runtime_class = vars(module).get(class_name) + if ( + isinstance(runtime_class, type) + and issubclass(runtime_class, pydantic.BaseModel) + and runtime_class.__module__ == module.__name__ + ): + models.append((module_class, runtime_class)) + return tuple(models) + + +def _unresolved_annotation(model: type[pydantic.BaseModel]) -> str | None: + try: + model.model_rebuild(raise_errors=True) + except pydantic.PydanticUndefinedAnnotation as error: + return str(error) + return None + + +def test_every_model_is_runtime_evaluated_for_ruff() -> None: + listed = _listed_runtime_evaluated_bases() + models = _model_classes() + missing = [ + f"{module_class.qualified_name}: add {_bases_to_list(module_class)}" + for module_class, _ in models + if not is_runtime_evaluated(module_class, listed) + ] + + assert models, "found no Pydantic models to check" + assert missing == [], ( + "ruff would not treat these Pydantic models as runtime-evaluated, so it " + "would move imports their annotations need under TYPE_CHECKING. Ruff " + "reads a base through the imports of the model's own module, so add the " + "base path named to runtime-evaluated-base-classes in ruff.toml, or " + "import the base from a path listed there:\n" + "\n".join(missing) + ) + + +def test_every_model_resolves_its_annotations() -> None: + unresolved = [ + f"{module_class.qualified_name}: {message}" + for module_class, model in _model_classes() + if (message := _unresolved_annotation(model)) is not None + ] + + assert unresolved == [], ( + "these models use names in annotations that don't exist at runtime, " + "usually an import moved under TYPE_CHECKING for a model ruff does not " + "treat as runtime-evaluated:\n" + "\n".join(unresolved) + ) + + +RESOLUTION_LISTED = frozenset( + {"pkg.models.Base", "pydantic.BaseModel", "pydantic.RootModel"} +) + + +@pytest.mark.parametrize( + ("source", "expected_paths", "runtime_evaluated"), + [ + pytest.param( + "from pydantic import BaseModel\n\nclass Model(BaseModel): ...", + ["pydantic.BaseModel"], + True, + id="listed-base", + ), + pytest.param( + "from pydantic import RootModel\n\nclass Model(RootModel[int]): ...", + ["pydantic.RootModel"], + True, + id="listed-base-subscripted", + ), + pytest.param( + "from ..models import Base\n\nclass Model(Base): ...", + ["pkg.models.Base"], + True, + id="listed-base-relative-import", + ), + pytest.param( + "from typing import Generic, TypeVar\n" + "from pydantic import BaseModel\n" + "Item = TypeVar('Item')\n" + "class _Root(BaseModel, Generic[Item]): ...\n" + "class _Middle(_Root[Item]): ...\n" + "class Model(_Middle[int]): ...", + [ + "pkg.service.api._Middle", + "pkg.service.api._Root", + "pydantic.BaseModel", + "typing.Generic", + ], + True, + id="same-module-chain", + ), + pytest.param( + "from pkg.reexports import Base\n\nclass Model(Base): ...", + ["pkg.reexports.Base"], + False, + id="unlisted-re-export", + ), + pytest.param( + "from pydantic_settings import BaseSettings\n\n" + "class Model(BaseSettings): ...", + ["pydantic_settings.BaseSettings"], + False, + id="unlisted-third-party-base", + ), + pytest.param( + "from pydantic import BaseModel\nBase = BaseModel\nclass Model(Base): ...", + [None], + False, + id="assigned-alias-unresolved", + ), + ], +) +def test_bases_resolve_as_ruff_resolves_them(source, expected_paths, runtime_evaluated): + model = module_classes("pkg.service.api", source, is_package=False)["Model"] + + assert [base.path for base in resolved_bases(model)] == expected_paths + assert is_runtime_evaluated(model, RESOLUTION_LISTED) is runtime_evaluated + + +def test_same_module_base_in_a_package_is_named_after_its_file() -> None: + """Ruff names a class in ``__init__.py`` with ``__init__`` in its path.""" + source = ( + "from somewhere import Unlisted\n\n" + "class Base(Unlisted): ...\n" + "class Model(Base): ...\n" + ) + model = module_classes("pkg.service", source, is_package=True)["Model"] + + assert is_runtime_evaluated(model, frozenset({"pkg.service.__init__.Base"})) + assert not is_runtime_evaluated(model, frozenset({"pkg.service.Base"})) diff --git a/projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py b/projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py new file mode 100644 index 000000000..b3ab7283c --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py @@ -0,0 +1,463 @@ +"""Pydantic models keep the runtime imports their annotations need. + +The packages' modules use ``from __future__ import annotations``, so ruff's +typing-only import rules (TC) move an import that only annotations use under +``if TYPE_CHECKING:``. Pydantic evaluates model annotations at runtime, so ruff +must treat every model as runtime-evaluated. It does when a base of the model, +read through the imports of the model's own module, is listed in +``runtime-evaluated-base-classes`` in ruff.toml, or is a class defined in that +same module that is runtime-evaluated itself. Ruff never follows a re-export to +the module that defines a class, so these checks resolve bases exactly as ruff +does. See docs/engineering/skills/typing-only-imports.md. + +The entrypoint's test suite keeps a copy of the resolution helpers, because +the two projects share no package that test helpers could live in. +""" + +from __future__ import annotations + +import ast +import builtins +import dataclasses +import importlib +import importlib.util +import pkgutil +import sys +from functools import cache +from pathlib import Path +from typing import TYPE_CHECKING + +import pydantic +import pytest +import tomllib + +if TYPE_CHECKING: + from collections.abc import Iterator + from types import ModuleType + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +SOURCE_ROOT = Path(__file__).resolve().parents[1] / "src" +# The executor and every project library installed in its environment. +CHECKED_PACKAGES = ( + "policyengine_fastapi", + "policyengine_simulation_contract", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_stage12_persistence", +) +# src/modal holds the Modal entry files. They are not part of the installed +# package, and importing them builds Modal images and functions, which the tests +# that cover them do only with a fake Modal SDK and a patched environment. They +# are checked statically instead, by test_modal_entry_files_define_no_models. +MODAL_ENTRY_DIRECTORY = SOURCE_ROOT / "modal" + +# What a module-level name is bound to: the qualified name an import gives it, +# a class statement in the same module, or None when ruff cannot resolve it (an +# assignment or a function definition). +Binding = str | ast.ClassDef | None + + +@dataclasses.dataclass(frozen=True) +class ModuleClass: + """A module-level class statement, with the names bound where it runs.""" + + module_name: str + ruff_module_path: str + node: ast.ClassDef + bindings: dict[str, Binding] + + @property + def qualified_name(self) -> str: + return f"{self.module_name}.{self.node.name}" + + +@dataclasses.dataclass(frozen=True) +class ResolvedBase: + """A base as ruff reads it, and whether ruff follows it within the module.""" + + written: str + path: str | None + same_module: bool + + +def _module_level_statements(body: list[ast.stmt]) -> Iterator[ast.stmt]: + """Yield the statements that run at module level, in source order.""" + for statement in body: + yield statement + if isinstance(statement, ast.If): + yield from _module_level_statements(statement.body) + yield from _module_level_statements(statement.orelse) + elif isinstance(statement, (ast.Try, ast.TryStar)): + for block in ( + statement.body, + *(handler.body for handler in statement.handlers), + statement.orelse, + statement.finalbody, + ): + yield from _module_level_statements(block) + + +def _names_bound_by(statement: ast.stmt, package: str) -> dict[str, Binding]: + """Return the module-level names a statement binds, as ruff resolves them.""" + if isinstance(statement, ast.Import): + bound: dict[str, Binding] = {} + for alias in statement.names: + if alias.asname is None: + top_level = alias.name.partition(".")[0] + bound[top_level] = top_level + else: + bound[alias.asname] = alias.name + return bound + if isinstance(statement, ast.ImportFrom): + source_module = ( + importlib.util.resolve_name( + "." * statement.level + (statement.module or ""), package + ) + if statement.level + else statement.module + ) + return { + alias.asname or alias.name: f"{source_module}.{alias.name}" + for alias in statement.names + if alias.name != "*" + } + if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef)): + return {statement.name: None} + if isinstance(statement, ast.Assign) or ( + isinstance(statement, ast.AnnAssign) and statement.value is not None + ): + targets = ( + statement.targets + if isinstance(statement, ast.Assign) + else [statement.target] + ) + return { + node.id: None + for target in targets + for node in ast.walk(target) + if isinstance(node, ast.Name) + } + return {} + + +def module_classes( + module_name: str, source: str, *, is_package: bool +) -> dict[str, ModuleClass]: + """Return each module-level class with the names bound at its statement. + + Ruff names a module after its file, so a package's ``__init__.py`` is + ``package.__init__`` to ruff. + """ + package = module_name if is_package else module_name.rpartition(".")[0] + ruff_module_path = f"{module_name}.__init__" if is_package else module_name + bindings: dict[str, Binding] = {} + classes: dict[str, ModuleClass] = {} + for statement in _module_level_statements(ast.parse(source).body): + if isinstance(statement, ast.ClassDef): + classes[statement.name] = ModuleClass( + module_name, ruff_module_path, statement, dict(bindings) + ) + bindings[statement.name] = statement + else: + bindings.update(_names_bound_by(statement, package)) + return classes + + +def _dotted_name(base: ast.expr) -> list[str] | None: + """Split ``a.b.Base`` into its names, reading ``Base[Item]`` as ``Base``.""" + if isinstance(base, ast.Subscript): + base = base.value + attributes: list[str] = [] + while isinstance(base, ast.Attribute): + attributes.insert(0, base.attr) + base = base.value + if not isinstance(base, ast.Name): + return None + return [base.id, *attributes] + + +def _ruff_path(module_class: ModuleClass, base: ast.expr) -> str | None: + """Return the qualified name ruff resolves a base to, through the imports.""" + names = _dotted_name(base) + if names is None: + return None + head, *attributes = names + if head not in module_class.bindings: + return ".".join(["builtins", *names]) if hasattr(builtins, head) else None + binding = module_class.bindings[head] + if binding is None: + return None + if isinstance(binding, ast.ClassDef): + binding = f"{module_class.ruff_module_path}.{head}" + return ".".join([binding, *attributes]) + + +def resolved_bases( + module_class: ModuleClass, followed: frozenset[str] = frozenset() +) -> list[ResolvedBase]: + """Return the bases ruff checks, following classes in the same module. + + Ruff resolves a same-module base's own bases with the names bound at the + subclass, so the subclass's bindings are kept on the way down. + """ + found: list[ResolvedBase] = [] + for base in module_class.node.bases: + names = _dotted_name(base) + local_class = ( + module_class.bindings.get(names[0]) + if names is not None and len(names) == 1 + else None + ) + same_module = isinstance(local_class, ast.ClassDef) + found.append( + ResolvedBase(ast.unparse(base), _ruff_path(module_class, base), same_module) + ) + if isinstance(local_class, ast.ClassDef) and local_class.name not in followed: + found += resolved_bases( + dataclasses.replace(module_class, node=local_class), + followed | {local_class.name}, + ) + return found + + +def is_runtime_evaluated(module_class: ModuleClass, listed: frozenset[str]) -> bool: + """Whether ruff treats the class's annotations as evaluated at runtime.""" + return any(base.path in listed for base in resolved_bases(module_class)) + + +def _names_a_model(path: str) -> bool: + try: + target = pkgutil.resolve_name(path) + except (AttributeError, ImportError, ValueError): + return False + return isinstance(target, type) and issubclass(target, pydantic.BaseModel) + + +def _bases_to_list(module_class: ModuleClass) -> list[str]: + """Return the model bases from other modules that would fix the class.""" + outside = [ + base.path or base.written + for base in resolved_bases(module_class) + if not base.same_module + ] + return [path for path in outside if _names_a_model(path)] or outside + + +def _listed_runtime_evaluated_bases() -> frozenset[str]: + with (REPOSITORY_ROOT / "ruff.toml").open("rb") as ruff_config: + settings = tomllib.load(ruff_config) + type_checking = settings["lint"]["flake8-type-checking"] + return frozenset(type_checking["runtime-evaluated-base-classes"]) + + +@cache +def _imported_modules() -> tuple[ModuleType, ...]: + modules = [] + for package_name in CHECKED_PACKAGES: + package = importlib.import_module(package_name) + modules.append(package) + modules += [ + importlib.import_module(module_info.name) + for module_info in pkgutil.walk_packages( + package.__path__, f"{package_name}." + ) + ] + return tuple(modules) + + +@cache +def _model_classes() -> tuple[tuple[ModuleClass, type[pydantic.BaseModel]], ...]: + """Return each Pydantic model a module-level class statement defines.""" + models = [] + for module in _imported_modules(): + source_path = Path(module.__file__) + classes = module_classes( + module.__name__, + source_path.read_text(encoding="utf-8"), + is_package=source_path.name == "__init__.py", + ) + for class_name, module_class in classes.items(): + runtime_class = vars(module).get(class_name) + if ( + isinstance(runtime_class, type) + and issubclass(runtime_class, pydantic.BaseModel) + and runtime_class.__module__ == module.__name__ + ): + models.append((module_class, runtime_class)) + return tuple(models) + + +def _unresolved_annotation(model: type[pydantic.BaseModel]) -> str | None: + try: + model.model_rebuild(raise_errors=True) + except pydantic.PydanticUndefinedAnnotation as error: + return str(error) + return None + + +def test_every_model_is_runtime_evaluated_for_ruff() -> None: + listed = _listed_runtime_evaluated_bases() + models = _model_classes() + missing = [ + f"{module_class.qualified_name}: add {_bases_to_list(module_class)}" + for module_class, _ in models + if not is_runtime_evaluated(module_class, listed) + ] + + assert models, "found no Pydantic models to check" + assert missing == [], ( + "ruff would not treat these Pydantic models as runtime-evaluated, so it " + "would move imports their annotations need under TYPE_CHECKING. Ruff " + "reads a base through the imports of the model's own module, so add the " + "base path named to runtime-evaluated-base-classes in ruff.toml, or " + "import the base from a path listed there:\n" + "\n".join(missing) + ) + + +def test_every_model_resolves_its_annotations() -> None: + unresolved = [ + f"{module_class.qualified_name}: {message}" + for module_class, model in _model_classes() + if (message := _unresolved_annotation(model)) is not None + ] + + assert unresolved == [], ( + "these models use names in annotations that don't exist at runtime, " + "usually an import moved under TYPE_CHECKING for a model ruff does not " + "treat as runtime-evaluated:\n" + "\n".join(unresolved) + ) + + +def test_modal_entry_files_define_no_models() -> None: + """Every class in src/modal derives only from builtins and the standard library.""" + standard_library = sys.stdlib_module_names + possible_models = [] + for source_path in sorted(MODAL_ENTRY_DIRECTORY.rglob("*.py")): + module_path = source_path.relative_to(SOURCE_ROOT).with_suffix("").parts + is_package = module_path[-1] == "__init__" + classes = module_classes( + ".".join(module_path[:-1] if is_package else module_path), + source_path.read_text(encoding="utf-8"), + is_package=is_package, + ) + possible_models += [ + f"{source_path.relative_to(REPOSITORY_ROOT)}: {module_class.node.name}" + f" derives from {base.path or base.written}" + for module_class in classes.values() + for base in resolved_bases(module_class) + if not base.same_module + and (base.path or "").partition(".")[0] not in standard_library + ] + + assert possible_models == [], ( + "these checks do not import src/modal, so a Pydantic model there would " + "go unchecked; define it in policyengine_simulation_executor instead:\n" + + "\n".join(possible_models) + ) + + +def test_rule_proof_runtime_evaluated_models_resolve() -> None: + """Proof for rule sim.quality.runtime-evaluated-models-resolve. + + The rule is registered in rules/registry.toml. + """ + fixtures = ( + REPOSITORY_ROOT / "rules/fixtures/sim.quality.runtime-evaluated-models-resolve" + ) + listed = _listed_runtime_evaluated_bases() + verdicts = { + fixture_name: is_runtime_evaluated( + module_classes( + "policyengine_simulation_executor.compat_models", + (fixtures / fixture_name).read_text(encoding="utf-8"), + is_package=False, + )["EconomyComparison"], + listed, + ) + for fixture_name in ("bad_compat_models.py", "good_compat_models.py") + } + + assert verdicts == {"bad_compat_models.py": False, "good_compat_models.py": True} + + +RESOLUTION_LISTED = frozenset( + {"pkg.models.Base", "pydantic.BaseModel", "pydantic.RootModel"} +) + + +@pytest.mark.parametrize( + ("source", "expected_paths", "runtime_evaluated"), + [ + pytest.param( + "from pydantic import BaseModel\n\nclass Model(BaseModel): ...", + ["pydantic.BaseModel"], + True, + id="listed-base", + ), + pytest.param( + "from pydantic import RootModel\n\nclass Model(RootModel[int]): ...", + ["pydantic.RootModel"], + True, + id="listed-base-subscripted", + ), + pytest.param( + "from ..models import Base\n\nclass Model(Base): ...", + ["pkg.models.Base"], + True, + id="listed-base-relative-import", + ), + pytest.param( + "from typing import Generic, TypeVar\n" + "from pydantic import BaseModel\n" + "Item = TypeVar('Item')\n" + "class _Root(BaseModel, Generic[Item]): ...\n" + "class _Middle(_Root[Item]): ...\n" + "class Model(_Middle[int]): ...", + [ + "pkg.service.api._Middle", + "pkg.service.api._Root", + "pydantic.BaseModel", + "typing.Generic", + ], + True, + id="same-module-chain", + ), + pytest.param( + "from pkg.reexports import Base\n\nclass Model(Base): ...", + ["pkg.reexports.Base"], + False, + id="unlisted-re-export", + ), + pytest.param( + "from pydantic_settings import BaseSettings\n\n" + "class Model(BaseSettings): ...", + ["pydantic_settings.BaseSettings"], + False, + id="unlisted-third-party-base", + ), + pytest.param( + "from pydantic import BaseModel\nBase = BaseModel\nclass Model(Base): ...", + [None], + False, + id="assigned-alias-unresolved", + ), + ], +) +def test_bases_resolve_as_ruff_resolves_them(source, expected_paths, runtime_evaluated): + model = module_classes("pkg.service.api", source, is_package=False)["Model"] + + assert [base.path for base in resolved_bases(model)] == expected_paths + assert is_runtime_evaluated(model, RESOLUTION_LISTED) is runtime_evaluated + + +def test_same_module_base_in_a_package_is_named_after_its_file() -> None: + """Ruff names a class in ``__init__.py`` with ``__init__`` in its path.""" + source = ( + "from somewhere import Unlisted\n\n" + "class Base(Unlisted): ...\n" + "class Model(Base): ...\n" + ) + model = module_classes("pkg.service", source, is_package=True)["Model"] + + assert is_runtime_evaluated(model, frozenset({"pkg.service.__init__.Base"})) + assert not is_runtime_evaluated(model, frozenset({"pkg.service.Base"})) diff --git a/rules/fixtures/sim.quality.runtime-evaluated-models-resolve/bad_compat_models.py b/rules/fixtures/sim.quality.runtime-evaluated-models-resolve/bad_compat_models.py new file mode 100644 index 000000000..1ed35d427 --- /dev/null +++ b/rules/fixtures/sim.quality.runtime-evaluated-models-resolve/bad_compat_models.py @@ -0,0 +1,16 @@ +"""Proof fixture: a model whose base ruff reads through a re-export. + +The rule must reject it. This is compat_models.py before the fix: ruff reads +the base as the executor's simulation_macro_output re-export, which is not +listed, so it would move imports the model's annotations need. +""" + +from __future__ import annotations + +from policyengine_simulation_executor.simulation_macro_output import ( + SingleYearMacroOutput, +) + + +class EconomyComparison(SingleYearMacroOutput): + """Legacy response schema name kept for generated clients.""" diff --git a/rules/fixtures/sim.quality.runtime-evaluated-models-resolve/good_compat_models.py b/rules/fixtures/sim.quality.runtime-evaluated-models-resolve/good_compat_models.py new file mode 100644 index 000000000..3faf99b2c --- /dev/null +++ b/rules/fixtures/sim.quality.runtime-evaluated-models-resolve/good_compat_models.py @@ -0,0 +1,13 @@ +"""Proof fixture: a model whose base ruff reads at a listed path. + +The rule must accept it. This is compat_models.py after the fix: the base is +imported from the contract module that defines it, the path ruff.toml lists. +""" + +from __future__ import annotations + +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput + + +class EconomyComparison(SingleYearMacroOutput): + """Legacy response schema name kept for generated clients.""" From 440088552b1d1f34c22350f898fb554a99c7771d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:51:58 +0400 Subject: [PATCH 08/10] Register the runtime-evaluated model rule and document typing-only imports sim.quality.runtime-evaluated-models-resolve is enforced by the two new test files, which run in the executor and entrypoint jobs of the PR test matrix. Its nodes span both projects, so they are listed from the repository root without a project. The test files join the guard-test discovery globs, so any test added to them must be registered. docs/engineering/skills/typing-only-imports.md explains when ruff moves an import, how it decides a model is runtime-evaluated, and how to list model bases; AGENTS.md and the skills README point to it. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 5 + changelog_entry.yaml | 1 + docs/engineering/skills/README.md | 3 + .../engineering/skills/typing-only-imports.md | 101 ++++++++++++++++++ docs/generated/rules.md | 11 +- rules/registry.toml | 32 ++++++ 6 files changed, 150 insertions(+), 3 deletions(-) create mode 100644 docs/engineering/skills/typing-only-imports.md diff --git a/AGENTS.md b/AGENTS.md index 62885a850..82a72aa1c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,6 +25,10 @@ When adding, changing, or removing a rule, a guard test, a checker configuration, or a CI job that enforces one, read `docs/engineering/skills/rules-registry.md`. +Before adding a Pydantic model or a class that other models inherit from, or +moving an import under `if TYPE_CHECKING:`, read +`docs/engineering/skills/typing-only-imports.md`. + ## Development - Use Python 3.13 and `uv`. @@ -140,6 +144,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. - `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.runtime-evaluated-models-resolve`: Every Pydantic model is runtime-evaluated for ruff: one of its bases, read through the imports of the model's own module, is listed in ruff.toml's runtime-evaluated-base-classes or is a runtime-evaluated class in that same module; and every Pydantic model resolves its annotations once its package is imported. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py`, `projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py`) - `sim.quality.tests-pass`: Each project's unit tests pass in its own locked environment. (blocking: CI job `test`) ### CI workflows and scripts diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e1109ee7e..1c1fb887c 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -5,6 +5,7 @@ - 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 typing-only import rules (TC) and moved annotation-only imports under TYPE_CHECKING, keeping runtime imports for Pydantic models and FastAPI endpoints + - Listed policyengine.py's Simulation, pydantic-settings' BaseSettings and the contract's SingleYearMacroOutput as runtime-evaluated bases for ruff, imported SingleYearMacroOutput from the contract in the executor's compat models instead of through a re-export, and added executor and entrypoint tests that fail when ruff, reading a model's bases through its module's imports, would not treat a Pydantic model as runtime-evaluated, or when a model's annotations do not resolve - Ran ruff check in CI with the first group of rule families (imports, pyupgrade and other safe auto-fixes) and fixed every finding - 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 - Moved multi-line `run` blocks in the PR, image smoke and client publishing workflows into single-line steps and tested scripts, tested the dependency-update scripts, fixed the shellcheck finding in modal-precompute.sh, and enforced the CI-script testing and shellcheck rules with empty baselines diff --git a/docs/engineering/skills/README.md b/docs/engineering/skills/README.md index f3f02b96c..33f2bd188 100644 --- a/docs/engineering/skills/README.md +++ b/docs/engineering/skills/README.md @@ -20,3 +20,6 @@ Current skills: isolation. - `rules-registry.md`: how rules are registered in `rules/registry.toml`, enforced, proven, and rendered into `docs/generated/rules.md`. +- `typing-only-imports.md`: which imports ruff moves under `TYPE_CHECKING`, how + it decides a Pydantic model is runtime-evaluated, and how to list model bases + so models keep the imports their annotations need. diff --git a/docs/engineering/skills/typing-only-imports.md b/docs/engineering/skills/typing-only-imports.md new file mode 100644 index 000000000..c20172ed2 --- /dev/null +++ b/docs/engineering/skills/typing-only-imports.md @@ -0,0 +1,101 @@ +# Typing-Only Imports Skill + +Read this before adding a Pydantic model, before adding a class that other +models inherit from, or before moving an import under `if TYPE_CHECKING:`. + +## When an import moves + +The package modules use `from __future__ import annotations`, so Python stores +every annotation as a string and never evaluates it. An import that only +annotations use is therefore not needed at runtime, and ruff's typing-only +import rules (`TC`) move it under `if TYPE_CHECKING:`, which is `False` when the +code runs and `True` for type checkers. That removes import-time work and +import cycles without changing behaviour. + +## Models that read annotations at runtime + +Some code does evaluate those strings at runtime: + +- Pydantic builds each model's validators from its annotations, including + `pydantic-settings` settings classes. +- FastAPI reads endpoint signatures when routes are registered. + +Names those annotations use must exist at runtime. `[lint.flake8-type-checking]` +in `ruff.toml` tells ruff which classes and functions these are: +`runtime-evaluated-base-classes` lists model base classes, and +`runtime-evaluated-decorators` lists the FastAPI route decorators. Ruff keeps an +import at runtime when a runtime-evaluated class or endpoint uses it in an +annotation. + +If ruff does not treat a model as runtime-evaluated, it moves the model's +annotation-only imports under `TYPE_CHECKING`, and Pydantic raises "not fully +defined" the first time the model is used. + +## How ruff decides a class is runtime-evaluated + +Ruff reads each base of a class statement through the import statements of the +module that holds the class, and compares the qualified name it gets with +`runtime-evaluated-base-classes`: + +- `from pydantic import BaseModel` makes the base `BaseModel` read as + `pydantic.BaseModel`, and so does `pydantic.BaseModel` after + `import pydantic`. A relative import resolves against the module's package. +- A subscripted base reads as the class it subscripts: `RootModel[Item]` is + `pydantic.RootModel` and `Generic[Item]` is `typing.Generic`. +- A base defined in the same module counts when it is runtime-evaluated itself. + Ruff follows inheritance through classes in the same module, however many + levels there are. +- Ruff follows nothing else. It never looks inside the module an imported base + comes from, so a base imported through a re-export reads as the re-exporting + path, and a model that subclasses a model from another module is covered only + when that base's import path is listed. A name bound by an assignment, such as + `Base = BaseModel`, does not resolve at all. + +Ruff does not check whether a class is a Pydantic model: listing the right path +is the only signal it gets. + +## Adding a model or a model base + +- Subclass a project model by importing it from the module that defines it, and + list that path, the defining module then the class name. Don't subclass + through a re-export such as the executor's `simulation_macro_output` + compatibility module: the re-exported path is not listed. +- When you add a model class that models in other modules inherit from, add its + path to `runtime-evaluated-base-classes` in the same change. +- A third-party base is listed by the public path its subclasses import it from, + such as `pydantic_settings.BaseSettings` or policyengine.py's + `policyengine.core.Simulation`. +- A base used only by models in its own module needs no entry, as long as it + derives from a listed class. + +If ruff reports a typing-only import in a module that defines a model, check the +model is runtime-evaluated before moving the import. Moving it for a model ruff +misses breaks the model at runtime. + +## The guard + +`sim.quality.runtime-evaluated-models-resolve` runs +`tests/test_runtime_evaluated_annotations.py` in two test suites: + +- the executor's, over the executor package and the project libraries installed + with it (`policyengine_fastapi`, `policyengine_simulation_contract`, + `policyengine_simulation_observability` and + `policyengine_stage12_persistence`); +- the entrypoint's, over `policyengine_simulation_entry`. + +Each imports every module of its packages and checks every Pydantic model a +module-level class statement defines: + +- the model is runtime-evaluated, with its bases resolved as ruff resolves them + above; a failure names the model and the base path to list; +- the model resolves its annotations, with `model_rebuild(raise_errors=True)`. + +The gateway defines no Pydantic models; it uses the contract's. The executor's +`src/modal` directory holds Modal entry files, which are not part of the +installed package and build Modal images and functions when imported, so the +guard does not import them; it checks statically that every class there derives +only from builtins and the standard library, so none can be a Pydantic model. + +The two projects share no package that test helpers could live in, so the +entrypoint's test keeps a copy of the executor's resolution helpers. Change both +copies together. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index ed76deec6..65f0643c8 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 | 33 | 1 | 0 | 14 | 48 | +| mechanical | 34 | 1 | 0 | 14 | 49 | | metric | 0 | 0 | 0 | 0 | 0 | | judgment | 0 | 0 | 0 | 0 | 0 | -**Blocking:** 34 of 48 rules. +**Blocking:** 35 of 49 rules. ## Gaps @@ -250,6 +250,7 @@ flowchart LR | `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.runtime-evaluated-models-resolve` | Every Pydantic model is runtime-evaluated for ruff: one of its bases, read through the imports of the model's own module, is listed in ruff.toml's runtime-evaluated-base-classes or is a runtime-evaluated class in that same module; and every Pydantic model resolves its annotations once its package is imported. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py`, `projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py` | `pr.yml:test` | yes | | `sim.quality.tests-pass` | Each project's unit tests pass in its own locked environment. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — | ```mermaid @@ -262,15 +263,19 @@ flowchart LR 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)"] + r_sim_quality_runtime_evaluated_models_resolve("sim.quality.runtime-evaluated-models-resolve (enforced)") + e_pytest_projects_policyengine_simulation_entry_tests_test_runtime_evaluated_annotations_py__projects_policyengine_simulation_executor_tests_test_runtime_evaluated_annotations_py["pytest projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py, projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py"] + j_CI_pr_yml_test["CI pr.yml:test"] r_sim_quality_tests_pass("sim.quality.tests-pass (enforced)") e_CI_job_test["CI job test"] - j_CI_pr_yml_test["CI pr.yml:test"] 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_runtime_evaluated_models_resolve --> e_pytest_projects_policyengine_simulation_entry_tests_test_runtime_evaluated_annotations_py__projects_policyengine_simulation_executor_tests_test_runtime_evaluated_annotations_py + e_pytest_projects_policyengine_simulation_entry_tests_test_runtime_evaluated_annotations_py__projects_policyengine_simulation_executor_tests_test_runtime_evaluated_annotations_py --> j_CI_pr_yml_test r_sim_quality_tests_pass --> e_CI_job_test e_CI_job_test --> j_CI_pr_yml_test ``` diff --git a/rules/registry.toml b/rules/registry.toml index 93075414e..d0253b723 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -51,6 +51,7 @@ guard_test_globs = [ "projects/*/tests/test_*schema_authority*.py", "projects/policyengine-simulation-entry/tests/test_deployment_assets.py", "projects/policyengine-simulation-gateway/tests/test_openapi_golden.py", + "projects/*/tests/test_runtime_evaluated_annotations.py", ] ignore_tests = [ { node = "projects/policyengine-simulation-entry/tests/test_deployment_assets.py::test_the_expression_evaluator_follows_github_semantics", reason = "Checks the GitHub Actions expression evaluator that the workflow-condition tests use; it enforces no repository rule itself." }, @@ -766,6 +767,37 @@ autofix = false [rule.enforcer] kind = "none" +[[rule]] +id = "sim.quality.runtime-evaluated-models-resolve" +category = "quality" +statement = "Every Pydantic model is runtime-evaluated for ruff: one of its bases, read through the imports of the model's own module, is listed in ruff.toml's runtime-evaluated-base-classes or is a runtime-evaluated class in that same module; and every Pydantic model resolves its annotations once its package is imported." +rationale = "docs/engineering/skills/typing-only-imports.md#how-ruff-decides-a-class-is-runtime-evaluated" +scope = ["projects/*/src/**/*.py", "libs/*/src/**/*.py", "ruff.toml"] +tier = "mechanical" +status = "enforced" +autofix = false +[rule.enforcer] +kind = "pytest" +nodes = [ + "projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py::test_every_model_is_runtime_evaluated_for_ruff", + "projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py::test_every_model_resolves_its_annotations", + "projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py::test_modal_entry_files_define_no_models", + "projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py::test_bases_resolve_as_ruff_resolves_them", + "projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py::test_same_module_base_in_a_package_is_named_after_its_file", + "projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py::test_every_model_is_runtime_evaluated_for_ruff", + "projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py::test_every_model_resolves_its_annotations", + "projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py::test_bases_resolve_as_ruff_resolves_them", + "projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py::test_same_module_base_in_a_package_is_named_after_its_file", +] +[[rule.ci]] +workflow = ".github/workflows/pr.yml" +job = "test" +invocation = "uv run pytest tests/" +[rule.proof] +bad_fixture = "rules/fixtures/sim.quality.runtime-evaluated-models-resolve/bad_compat_models.py" +good_fixture = "rules/fixtures/sim.quality.runtime-evaluated-models-resolve/good_compat_models.py" +test = "projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py::test_rule_proof_runtime_evaluated_models_resolve" + # --- CI workflows and scripts ---------------------------------------------- [[rule]] From 5ee0b46ee47ab4d6356760e3cfe7b44f2de3e38d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:44:19 +0400 Subject: [PATCH 09/10] Move the typing-only imports #781 added under TYPE_CHECKING #781 landed on main after this stack branched. TC, which this PR enables, flags the imports its contract job_progress module uses only in annotations (Mapping and datetime), which move under TYPE_CHECKING, and the cast in its new integration poll test, whose Client type is quoted. The module's Pydantic models annotate their timestamps with AwareDatetime, which stays a runtime import. Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_contract/job_progress.py | 8 +++++--- .../tests/simulation/test_calculate.py | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py index cf362ecb9..71e8c2ed9 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py @@ -13,11 +13,9 @@ from __future__ import annotations -from collections.abc import Mapping -from datetime import datetime from enum import StrEnum from types import MappingProxyType -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal from pydantic import ( AwareDatetime, @@ -29,6 +27,10 @@ model_validator, ) +if TYPE_CHECKING: + from collections.abc import Mapping + from datetime import datetime + JOB_PROGRESS_CAPABILITY_KEY = "job_progress" JOB_PROGRESS_CAPABILITY_VERSION = 1 diff --git a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py index 726d5a741..3836602bf 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py @@ -436,7 +436,7 @@ def test_poll_for_completion_records_running_and_complete_bodies( monkeypatch.setattr(time, "sleep", lambda seconds: None) poll_for_completion( - cast(Client, object()), + cast("Client", object()), "fc-test", max_wait_seconds=1.0, poll_interval=0.25, From 8c3125d85173230caeca454c282a701b1577cc0e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:44:29 +0400 Subject: [PATCH 10/10] List the StageReportingSimulation base #781 added as runtime-evaluated #781 landed on main after this stack branched. It added StageReportingSimulation, a policyengine.core.Simulation subclass, and made ArtifactBaselineSimulation subclass it from another module. Ruff reads that base through baseline_artifacts' import, which is not listed, so it would no longer treat ArtifactBaselineSimulation as runtime-evaluated. The sim.quality.runtime-evaluated-models-resolve guard names the path to list; it is added beside the other project bases. Co-Authored-By: Claude Opus 5.5 --- ruff.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/ruff.toml b/ruff.toml index 88b366afc..a436b0904 100644 --- a/ruff.toml +++ b/ruff.toml @@ -36,6 +36,7 @@ runtime-evaluated-base-classes = [ "policyengine_simulation_contract.stage12_execution.StrictContractModel", "policyengine_simulation_contract.stage12_manifest.StrictManifestModel", "policyengine_simulation_executor.precompute_models._StrictModel", + "policyengine_simulation_executor.stage_reporting_simulation.StageReportingSimulation", "pydantic.BaseModel", "pydantic.RootModel", "pydantic_settings.BaseSettings",