diff --git a/AGENTS.md b/AGENTS.md index 26122fae6..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`. @@ -139,7 +143,8 @@ 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.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 83047a7a5..1c1fb887c 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -4,6 +4,8 @@ - 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 + - 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 620a8e7a6..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 @@ -249,7 +249,8 @@ 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.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 @@ -261,16 +262,20 @@ 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_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__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_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/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/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/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..3836602bf 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, @@ -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, 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-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/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..94b3d5ef3 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -25,11 +25,9 @@ import logging 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 +54,11 @@ import modal +if TYPE_CHECKING: + from collections.abc import Callable + + 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/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/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/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/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/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 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__) 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_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/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..a436b0904 100644 --- a/ruff.toml +++ b/ruff.toml @@ -11,4 +11,48 @@ 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 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", + "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", +] +# 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/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.""" diff --git a/rules/registry.toml b/rules/registry.toml index 605f0fa22..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." }, @@ -732,7 +733,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" @@ -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]] 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.