diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index d0d75e5e2..c69cb49af 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -99,6 +99,9 @@ jobs: - name: Run ruff format check run: uv run ruff format --check projects/*/src libs/*/src + - name: Run ruff check + run: uv run ruff check . + quality-guards: name: Quality guards # Full CI runs only on the lowest unmerged pull request in a stack. diff --git a/AGENTS.md b/AGENTS.md index 5a457a752..d815f7c91 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,7 +135,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. - `sim.quality.every-project-tested`: Every project and library with unit tests runs them in the PR test matrix. (doc-only, not checked: follow it yourself) - `sim.quality.ruff-format-src`: Source code under projects/*/src and libs/*/src is formatted with ruff format. (blocking: ruff format) - `sim.quality.ruff-format-tests`: Test code is formatted with ruff format too; CI currently checks only src directories. (doc-only, not checked: follow it yourself) -- `sim.quality.ruff-lint`: Python code passes ruff check with the rules selected in ruff.toml (ruff's defaults, made explicit); CI does not run ruff check yet. (advisory: ruff check (E4, E7, E9, F), not blocking) +- `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.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 a73b4fa61..ab91aaa8a 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -4,6 +4,7 @@ - Made CLAUDE.md and the Copilot instructions thin pointers to AGENTS.md, required workflow logic longer than one line to live in tested and shellchecked scripts, and registered the tracking-issue rule - Added a rules registry (rules/registry.toml) with a standard-library tool that validates it, checks every rule is enforced where it claims, and renders docs/generated/rules.md and the Rules block in AGENTS.md changed: + - 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, and enforced the one-line `run` rule with an empty 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 - Added CODEOWNERS for the rule machinery, pinned CI's ruff to 0.15.22 with an explicit rule selection, and documented that agents never bypass branch rules diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 334f82353..3032e4089 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 | 29 | 0 | 1 | 14 | 44 | +| mechanical | 30 | 0 | 0 | 14 | 44 | | metric | 0 | 0 | 0 | 0 | 0 | | judgment | 0 | 0 | 0 | 0 | 0 | -**Blocking:** 29 of 44 rules. +**Blocking:** 30 of 44 rules. ## Gaps @@ -35,7 +35,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf ### Advisory (checked but not blocking) -- `sim.quality.ruff-lint`: Python code passes ruff check with the rules selected in ruff.toml (ruff's defaults, made explicit); CI does not run ruff check yet. +- None. ### Unproven (no proof that the checker rejects a bad example) @@ -233,7 +233,7 @@ flowchart LR | `sim.quality.every-project-tested` | Every project and library with unit tests runs them in the PR test matrix. | mechanical | doc-only | no | nothing | — | — | | `sim.quality.ruff-format-src` | Source code under projects/*/src and libs/*/src is formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — | | `sim.quality.ruff-format-tests` | Test code is formatted with ruff format too; CI currently checks only src directories. | mechanical | doc-only | no | nothing | — | — | -| `sim.quality.ruff-lint` | Python code passes ruff check with the rules selected in ruff.toml (ruff's defaults, made explicit); CI does not run ruff check yet. | mechanical | advisory | no | ruff check (E4, E7, E9, F) | — | — | +| `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.tests-pass` | Each project's unit tests pass in its own locked environment. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — | ```mermaid @@ -244,9 +244,8 @@ flowchart LR e_ruff_format["ruff format"] 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 (advisory)") - e_ruff_check__E4__E7__E9__F_["ruff check (E4, E7, E9, F)"] - local_only>"not run in CI"] + 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)"] 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"] @@ -254,8 +253,8 @@ flowchart LR r_sim_quality_ruff_format_src --> e_ruff_format e_ruff_format --> j_CI_pr_yml_lint r_sim_quality_ruff_format_tests -.-> gap_unenforced - r_sim_quality_ruff_lint --> e_ruff_check__E4__E7__E9__F_ - e_ruff_check__E4__E7__E9__F_ -.-> local_only + 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_tests_pass --> e_CI_job_test e_CI_job_test --> j_CI_pr_yml_test ``` diff --git a/libs/policyengine-fastapi/fixtures/ping/shared.py b/libs/policyengine-fastapi/fixtures/ping/shared.py index b7f480c13..525a46661 100644 --- a/libs/policyengine-fastapi/fixtures/ping/shared.py +++ b/libs/policyengine-fastapi/fixtures/ping/shared.py @@ -1,9 +1,8 @@ """Shared fixtures for ping endpoint tests.""" +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -import pytest - from policyengine_fastapi import ping from policyengine_fastapi.health import HealthRegistry diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py index 7cc751967..04b09fae9 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py @@ -4,8 +4,16 @@ from .observability import ( ObservabilityConfig as ObservabilityConfig, +) +from .observability import ( SimulationStage as SimulationStage, +) +from .observability import ( TracerCaptureMode as TracerCaptureMode, +) +from .observability import ( build_observability as build_observability, +) +from .observability import ( get_observability as get_observability, ) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/__init__.py index 468252342..89c4db22d 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/__init__.py @@ -1 +1,3 @@ from .jwt_decoder import JWTDecoder + +__all__ = ["JWTDecoder"] diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py index 89cbd8cdf..07331be90 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py @@ -1,8 +1,8 @@ -from typing import Any +import logging + +import jwt from fastapi import Depends, HTTPException, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer -import jwt -import logging # Based on the auth0 article: https://auth0.com/blog/build-and-secure-fastapi-server-with-auth0/ @@ -35,8 +35,7 @@ def __call__( if self.auto_error: LOG.info("Invalid auth token: Throwing 4XX error") raise err - else: - LOG.info("Ignoring invalid auth token") + LOG.info("Ignoring invalid auth token") return None def _validate(self, token: HTTPAuthorizationCredentials | None): diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/database.py b/libs/policyengine-fastapi/src/policyengine_fastapi/database.py index 8293c919f..e5569f372 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/database.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/database.py @@ -1,5 +1,6 @@ -from pathlib import Path -from typing import Any, Callable, Generator +from collections.abc import Callable, Generator +from typing import Any + from sqlalchemy import Engine, StaticPool, create_engine from sqlmodel import Session diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py b/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py index 6ec5c443b..e2821e764 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py @@ -4,11 +4,9 @@ """ import logging +from collections.abc import Callable from contextlib import asynccontextmanager -from functools import wraps -from typing import Any, Callable, TypeVar, cast - -from fastapi import FastAPI +from typing import Any, TypeVar log = logging.getLogger(__name__) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py index d816e0a3f..464826daf 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py @@ -1,11 +1,12 @@ -from typing import Callable, Optional -from pydantic import BaseModel, Field +from collections.abc import Callable + +from pydantic import BaseModel class ProbeStatus(BaseModel): name: str healthy: bool - message: Optional[str] = None + message: str | None = None class SystemStatus(BaseModel): diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py index c26d4e1fc..da530d37b 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py @@ -9,30 +9,52 @@ from .config import ( ObservabilityConfig as ObservabilityConfig, +) +from .config import ( parse_header_value_pairs as parse_header_value_pairs, ) from .contracts import ( SimulationCompositeTraceResponse as SimulationCompositeTraceResponse, +) +from .contracts import ( SimulationLifecycleEvent as SimulationLifecycleEvent, +) +from .contracts import ( SimulationRunSummary as SimulationRunSummary, +) +from .contracts import ( SimulationTimelineEntry as SimulationTimelineEntry, +) +from .contracts import ( TracerArtifactManifest as TracerArtifactManifest, +) +from .contracts import ( VersionStageMetricResponse as VersionStageMetricResponse, ) from .correlation import ( generate_observability_id as generate_observability_id, +) +from .correlation import ( stable_config_hash as stable_config_hash, ) from .emitters import ( - Observability as Observability, NoOpObservability as NoOpObservability, +) +from .emitters import ( NoOpSpan as NoOpSpan, ) +from .emitters import ( + Observability as Observability, +) from .provider import ( build_observability as build_observability, +) +from .provider import ( get_observability as get_observability, ) from .stages import ( SimulationStage as SimulationStage, +) +from .stages import ( TracerCaptureMode as TracerCaptureMode, ) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py index 9a2dc05e3..1dc46157e 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py @@ -1,7 +1,7 @@ from __future__ import annotations -from hashlib import sha256 import json +from hashlib import sha256 from typing import Any from uuid import uuid4 diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py index 97afa0eee..376aaf46d 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py @@ -1,9 +1,9 @@ from __future__ import annotations +from collections.abc import Mapping from contextlib import AbstractContextManager from dataclasses import dataclass, field -from typing import Protocol -from typing import Any, Mapping +from typing import Any, Protocol from .config import ObservabilityConfig from .contracts import SimulationLifecycleEvent, TracerArtifactManifest diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/__init__.py index 7889d5062..7a0d28bfb 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/__init__.py @@ -1,8 +1,10 @@ +from .console import export_ot_to_console as export_ot_to_console from .gcp import ( GCPLoggingInstrumentor as GCPLoggingInstrumentor, +) +from .gcp import ( export_ot_to_gcp as export_ot_to_gcp, ) -from .console import export_ot_to_console as export_ot_to_console from .instrumentor import ( FastAPIEnhancedInstrumenter as FastAPIEnhancedInstrumenter, ) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py index 34c0475e4..dc265d56e 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py @@ -1,22 +1,18 @@ +from opentelemetry import metrics, trace +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import ( + ConsoleMetricExporter, + PeriodicExportingMetricReader, +) from opentelemetry.sdk.resources import ( - SERVICE_NAME, - SERVICE_INSTANCE_ID, Resource, ) - -from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import ( BatchSpanProcessor, ConsoleSpanExporter, ) -from opentelemetry import metrics -from opentelemetry.sdk.metrics import MeterProvider -from opentelemetry.sdk.metrics.export import ( - PeriodicExportingMetricReader, - ConsoleMetricExporter, -) from policyengine_fastapi.exit import exit diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py index 442e3aad2..968fed712 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py @@ -1,14 +1,12 @@ import logging import os -from opentelemetry.sdk.resources import Resource -from opentelemetry import trace -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor - -from opentelemetry import metrics +from opentelemetry import metrics, trace from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor from policyengine_fastapi.exit import exit @@ -96,10 +94,10 @@ def instrument(self): def export_ot_to_gcp(resource: Resource): - from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter from opentelemetry.exporter.cloud_monitoring import ( CloudMonitoringMetricsExporter, ) + from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter """ configure opentelemetry to directly export to gcp cloudtrace/metrics diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py index 534cfebbf..c91265317 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py @@ -1,7 +1,7 @@ from fastapi import FastAPI +from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor from .middleware import Middleware -from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor class FastAPIEnhancedInstrumenter: diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py index 6702e76cd..e39ea23e0 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py @@ -1,6 +1,7 @@ from dataclasses import dataclass from time import time from typing import Any + from fastapi import FastAPI, Request from fastapi.routing import APIRoute from opentelemetry import metrics diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py index 877e93db8..bd53cb5ac 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py @@ -1,6 +1,7 @@ -from fastapi import FastAPI, APIRouter +from fastapi import APIRouter, FastAPI from fastapi.responses import JSONResponse from pydantic import BaseModel + from policyengine_fastapi.health import HealthRegistry, HealthStatus diff --git a/libs/policyengine-fastapi/tests/test_exit.py b/libs/policyengine-fastapi/tests/test_exit.py index c159a7b9a..dd580fe03 100644 --- a/libs/policyengine-fastapi/tests/test_exit.py +++ b/libs/policyengine-fastapi/tests/test_exit.py @@ -1,6 +1,7 @@ from unittest.mock import Mock, call -from policyengine_fastapi.exit import exit + import pytest +from policyengine_fastapi.exit import exit @pytest.mark.asyncio @@ -14,7 +15,7 @@ async def test_lifecycle_executes_callbacks__when_app_ends(): exit()(callback) - async with exit.lifespan() as resource: + async with exit.lifespan(): app() parent.assert_has_calls([call.app(), call.callback()]) diff --git a/libs/policyengine-fastapi/tests/test_observability.py b/libs/policyengine-fastapi/tests/test_observability.py index e40a35856..9c3513428 100644 --- a/libs/policyengine-fastapi/tests/test_observability.py +++ b/libs/policyengine-fastapi/tests/test_observability.py @@ -1,4 +1,4 @@ -from datetime import datetime, UTC +from datetime import UTC, datetime from policyengine_fastapi.observability import ( NoOpObservability, diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index 97c16a11b..533860436 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py @@ -3,8 +3,9 @@ """ import json -from typing import ClassVar, Literal, Optional +from typing import ClassVar, Literal +from policyengine_simulation_observability.telemetry import TelemetryEnvelope from pydantic import ( BaseModel, ConfigDict, @@ -14,15 +15,14 @@ model_validator, ) +from policyengine_simulation_contract.json_types import JsonObject +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput from policyengine_simulation_contract.spm import ( - SPMSelection, SPMCapability, SPMComparisonProvenance, SPMErrorDetail, + SPMSelection, ) -from policyengine_simulation_contract.json_types import JsonObject -from policyengine_simulation_contract.macro_output import SingleYearMacroOutput -from policyengine_simulation_observability.telemetry import TelemetryEnvelope # Hard cap on request body size (bytes). SimulationOptions + telemetry + any # reform/baseline parameter tree should fit comfortably in ~256 KB. A hostile @@ -100,24 +100,24 @@ class GatewayRequestBase(BaseModel): """ country: str - spm: Optional[SPMSelection] = None - version: Optional[str] = None - policyengine_version: Optional[str] = None + spm: SPMSelection | None = None + version: str | None = None + policyengine_version: str | None = None telemetry: TelemetryEnvelope | None = None # Fields forwarded to SimulationOptions on the worker side. - scope: Optional[str] = None - time_period: Optional[str] = None - reform: Optional[PolicyParameterChanges] = None - baseline: Optional[PolicyParameterChanges] = None - region_group: Optional[list[str]] = None - title: Optional[str] = None - include_cliffs: Optional[bool] = None - model_version: Optional[str] = None + scope: str | None = None + time_period: str | None = None + reform: PolicyParameterChanges | None = None + baseline: PolicyParameterChanges | None = None + region_group: list[str] | None = None + title: str | None = None + include_cliffs: bool | None = None + model_version: str | None = None # National macro requests run segmented across region groups by default; # False forces a monolithic single-process run. Ignored for requests that # are not eligible for segmentation (regional, household-scope, UK). - segmented: Optional[bool] = None + segmented: bool | None = None model_config = ConfigDict( extra="forbid", @@ -143,7 +143,7 @@ def enforce_max_payload_size(cls, value): class SimulationRequest(GatewayRequestBase): """Request model for simulation submission.""" - region: Optional[str] = None + region: str | None = None @model_validator(mode="after") def region_xor_region_group(self) -> "SimulationRequest": @@ -157,10 +157,10 @@ class PolicyEngineBundle(BaseModel): """Resolved runtime provenance returned by the gateway.""" model_version: str - policyengine_version: Optional[str] = None - data_version: Optional[str] = None - dataset: Optional[str] = None - spm: Optional[SPMCapability] = None + policyengine_version: str | None = None + data_version: str | None = None + dataset: str | None = None + spm: SPMCapability | None = None class SimulationErrorResponse(BaseModel): @@ -189,11 +189,11 @@ class JobStatusResponse(BaseModel): """Response model for job status polling.""" status: Literal["running", "complete", "failed"] - result: Optional[SingleYearMacroOutput] = None - errors: Optional[list[SPMErrorDetail]] = None - error: Optional[str] = None - resolved_app_name: Optional[str] = None - policyengine_bundle: Optional[PolicyEngineBundle] = None + result: SingleYearMacroOutput | None = None + errors: list[SPMErrorDetail] | None = None + error: str | None = None + resolved_app_name: str | None = None + policyengine_bundle: PolicyEngineBundle | None = None class BudgetWindowBatchRequest(GatewayRequestBase): @@ -239,7 +239,7 @@ class BudgetWindowAnnualImpact(BaseModel): model_config = ConfigDict(extra="forbid") - spm_provenance: Optional[SPMComparisonProvenance] = None + spm_provenance: SPMComparisonProvenance | None = None year: str taxRevenueImpact: float federalTaxRevenueImpact: float @@ -282,8 +282,8 @@ class BatchChildJobStatus(BaseModel): "failed", "cancelled", ] - errors: Optional[list[SPMErrorDetail]] = None - error: Optional[str] = None + errors: list[SPMErrorDetail] | None = None + error: str | None = None class BudgetWindowBatchSubmitResponse(BaseModel): @@ -302,17 +302,17 @@ class BudgetWindowBatchStatusResponse(BaseModel): """Response model for budget-window batch polling.""" status: Literal["submitted", "running", "complete", "failed"] - progress: Optional[int] = None + progress: int | None = None completed_years: list[str] = Field(default_factory=list) running_years: list[str] = Field(default_factory=list) queued_years: list[str] = Field(default_factory=list) failed_years: list[str] = Field(default_factory=list) child_jobs: dict[str, BatchChildJobStatus] = Field(default_factory=dict) - result: Optional[BudgetWindowResult] = None - errors: Optional[list[SPMErrorDetail]] = None - error: Optional[str] = None - resolved_app_name: Optional[str] = None - policyengine_bundle: Optional[PolicyEngineBundle] = None + result: BudgetWindowResult | None = None + errors: list[SPMErrorDetail] | None = None + error: str | None = None + resolved_app_name: str | None = None + policyengine_bundle: PolicyEngineBundle | None = None class BudgetWindowBatchState(BaseModel): @@ -339,12 +339,12 @@ class BudgetWindowBatchState(BaseModel): partial_annual_impacts: dict[str, BudgetWindowAnnualImpact] = Field( default_factory=dict ) - result: Optional[BudgetWindowResult] = None - errors: Optional[list[SPMErrorDetail]] = None - error: Optional[str] = None + result: BudgetWindowResult | None = None + errors: list[SPMErrorDetail] | None = None + error: str | None = None created_at: str updated_at: str - observability_id: Optional[str] = None + observability_id: str | None = None class PingRequest(BaseModel): diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/json_types.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/json_types.py index 1f6668bc6..26264046b 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/json_types.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/json_types.py @@ -4,5 +4,4 @@ from pydantic import JsonValue - type JsonObject = dict[str, JsonValue] diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index 1df9228a0..135bed7b8 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -3,7 +3,7 @@ import re from collections.abc import Sequence from datetime import date -from typing import Any, Literal, Optional +from typing import Any, Literal from pydantic import ( BaseModel, @@ -37,14 +37,12 @@ class SPMSelection(BaseModel): frozen=True, extra="forbid", json_schema_extra=_selection_schema ) - forecast_content_sha256: Optional[str] = Field( - default=None, pattern=r"^[0-9a-f]{64}$" - ) - scenario: Optional[str] = Field(default=None, min_length=1, pattern=r"^\S+$") + forecast_content_sha256: str | None = Field(default=None, pattern=r"^[0-9a-f]{64}$") + scenario: str | None = Field(default=None, min_length=1, pattern=r"^\S+$") geography_kind: Literal["county", "national", "metro"] = "county" - geography_id: Optional[str] = Field(default=None, min_length=1) + geography_id: str | None = Field(default=None, min_length=1) county_vintage: str = Field(default="2020", pattern=r"^[0-9]{4}$") - as_of: Optional[str] = None + as_of: str | None = None @model_serializer(mode="wrap") def serialize_selection( @@ -153,9 +151,9 @@ class SPMProvenance(BaseModel): forecast_sha256: str = Field(pattern=r"^[0-9a-f]{64}$") scenario: str = Field(min_length=1, pattern=r"^\S+$") geography_kind: Literal["county", "national", "metro"] - geography_id: Optional[str] = Field(min_length=1) + geography_id: str | None = Field(min_length=1) county_vintage: str = Field(pattern=r"^[0-9]{4}$") - as_of: Optional[date] + as_of: date | None years: tuple[str, ...] runtime_versions: "SPMRuntimeVersions" diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index 920aa1251..a048b8f3a 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -2,12 +2,12 @@ from __future__ import annotations +import json +import re from collections.abc import Mapping from copy import deepcopy from datetime import datetime from hashlib import sha256 -import json -import re from typing import Annotated, Literal, Protocol from packaging.version import InvalidVersion, Version diff --git a/libs/policyengine-simulation-contract/tests/conftest.py b/libs/policyengine-simulation-contract/tests/conftest.py index 1d30dfbde..36ff49a60 100644 --- a/libs/policyengine-simulation-contract/tests/conftest.py +++ b/libs/policyengine-simulation-contract/tests/conftest.py @@ -6,7 +6,6 @@ """ import pytest - from policyengine_simulation_contract import budget_window_state diff --git a/libs/policyengine-simulation-contract/tests/test_budget_window_state.py b/libs/policyengine-simulation-contract/tests/test_budget_window_state.py index beccd5c27..966e70b1c 100644 --- a/libs/policyengine-simulation-contract/tests/test_budget_window_state.py +++ b/libs/policyengine-simulation-contract/tests/test_budget_window_state.py @@ -5,8 +5,8 @@ create_initial_batch_state, get_batch_job_state, mark_batch_complete, - mark_batch_running, mark_batch_failed, + mark_batch_running, mark_child_completed, mark_child_failed, mark_child_started, diff --git a/libs/policyengine-simulation-contract/tests/test_dataset_uri.py b/libs/policyengine-simulation-contract/tests/test_dataset_uri.py index eaa6c6818..8a873e5e6 100644 --- a/libs/policyengine-simulation-contract/tests/test_dataset_uri.py +++ b/libs/policyengine-simulation-contract/tests/test_dataset_uri.py @@ -1,7 +1,6 @@ """Tests for dataset URI normalization.""" import pytest - from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri diff --git a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py index 1c44e1889..ab2473134 100644 --- a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py +++ b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py @@ -4,9 +4,8 @@ import json -import pytest - import policyengine_simulation_contract.hf_dataset as hf_dataset +import pytest from policyengine_simulation_contract.hf_dataset import ( HuggingFaceDatasetReferenceError, parse_hf_dataset_uri, diff --git a/libs/policyengine-simulation-contract/tests/test_spm_selection.py b/libs/policyengine-simulation-contract/tests/test_spm_selection.py index 0993b617f..3ad50ac18 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_selection.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_selection.py @@ -1,7 +1,6 @@ """Partial public selections inherit pinned bundle defaults before validation.""" import pytest - from policyengine_simulation_contract.gateway_models import SimulationRequest from policyengine_simulation_contract.spm import ( SPMInputError, diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_execution.py b/libs/policyengine-simulation-contract/tests/test_stage12_execution.py index 7c8a326a0..bc7c9eb84 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_execution.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_execution.py @@ -6,8 +6,6 @@ from uuid import UUID import pytest -from pydantic import JsonValue - from policyengine_simulation_contract.spm import ( SPMRuntimeVersions, SPMSelection, @@ -30,6 +28,7 @@ UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) +from pydantic import JsonValue def _plan() -> Stage12OutputPlan: diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 627a43f5b..ddf024aff 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -5,7 +5,6 @@ from datetime import datetime, timezone import pytest - from policyengine_simulation_contract.stage12_bundle import ( ResolvedStage12Bundle, Stage12BundleManifest, 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 d923c3379..6b8d6afbc 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,10 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from enum import StrEnum from types import MappingProxyType -from typing import Mapping class RunConfiguration(StrEnum): diff --git a/libs/policyengine-simulation-observability/tests/test_errors.py b/libs/policyengine-simulation-observability/tests/test_errors.py index 0b3301d4c..bca8a9266 100644 --- a/libs/policyengine-simulation-observability/tests/test_errors.py +++ b/libs/policyengine-simulation-observability/tests/test_errors.py @@ -7,7 +7,6 @@ from policyengine_simulation_observability import errors as errors_module - CORRELATION_RE = re.compile(r"correlation_id=([0-9a-f]{32})") diff --git a/libs/policyengine-simulation-observability/tests/test_observability.py b/libs/policyengine-simulation-observability/tests/test_observability.py index f0755284c..7892038dc 100644 --- a/libs/policyengine-simulation-observability/tests/test_observability.py +++ b/libs/policyengine-simulation-observability/tests/test_observability.py @@ -8,7 +8,6 @@ GoogleCloudLogDestination, StdoutLogDestination, ) - from policyengine_simulation_observability.identifiers import ( OBSERVABILITY_ID_HEADER, ) diff --git a/libs/policyengine-simulation-observability/tests/test_stages.py b/libs/policyengine-simulation-observability/tests/test_stages.py index 9481c909b..e3413410b 100644 --- a/libs/policyengine-simulation-observability/tests/test_stages.py +++ b/libs/policyengine-simulation-observability/tests/test_stages.py @@ -1,5 +1,4 @@ import pytest - from policyengine_simulation_observability.stages import ( RUN_STAGE_REGISTRY, RunConfiguration, diff --git a/libs/policyengine-simulation-observability/tests/test_telemetry.py b/libs/policyengine-simulation-observability/tests/test_telemetry.py index ff452d68c..39cfc5cac 100644 --- a/libs/policyengine-simulation-observability/tests/test_telemetry.py +++ b/libs/policyengine-simulation-observability/tests/test_telemetry.py @@ -1,5 +1,4 @@ import pytest - from policyengine_simulation_observability.telemetry import ( TelemetryEnvelope, normalize_observability_context, diff --git a/projects/policyengine-apis-integ/tests/simulation/conftest.py b/projects/policyengine-apis-integ/tests/simulation/conftest.py index 12fa39e22..c1e616b18 100644 --- a/projects/policyengine-apis-integ/tests/simulation/conftest.py +++ b/projects/policyengine-apis-integ/tests/simulation/conftest.py @@ -4,8 +4,6 @@ import httpx import pytest -from pydantic_settings import BaseSettings, SettingsConfigDict - from policyengine_api_simulation_client import AuthenticatedClient, Client from policyengine_api_simulation_client.api.default import ( get_budget_window_job_status_budget_window_jobs_batch_job_id_get, @@ -15,10 +13,10 @@ BudgetWindowBatchRequest, BudgetWindowBatchStatusResponse, ) +from pydantic_settings import BaseSettings, SettingsConfigDict from .polling import poll_retry_delay_seconds - BUDGET_WINDOW_YEARS = ["2026", "2027"] BUDGET_WINDOW_REFORM = { "gov.irs.credits.ctc.refundable.fully_refundable": {"2023-01-01.2100-12-31": True} diff --git a/projects/policyengine-apis-integ/tests/simulation/polling.py b/projects/policyengine-apis-integ/tests/simulation/polling.py index 091eaef93..d9c14c31a 100644 --- a/projects/policyengine-apis-integ/tests/simulation/polling.py +++ b/projects/policyengine-apis-integ/tests/simulation/polling.py @@ -4,7 +4,6 @@ from collections.abc import Mapping from http import HTTPStatus - TRANSIENT_POLL_STATUS_CODES: frozenset[HTTPStatus] = frozenset( { HTTPStatus.BAD_GATEWAY, diff --git a/projects/policyengine-apis-integ/tests/simulation/test_auth_smoke.py b/projects/policyengine-apis-integ/tests/simulation/test_auth_smoke.py index f8c3c770d..1b01bddcb 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_auth_smoke.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_auth_smoke.py @@ -7,7 +7,6 @@ from .conftest import settings - pytestmark = pytest.mark.skipif( not settings.gateway_auth_required, reason="Gateway auth is not required for this deployment; skipping auth smoke.", @@ -34,8 +33,7 @@ def test_gated_endpoint_rejects_missing_token() -> None: def test_gated_endpoint_accepts_valid_token() -> None: """A valid token must advance past auth and reach the job lookup.""" assert settings.access_token, ( - "simulation_integ_test_access_token must be set when gateway auth " - "is required" + "simulation_integ_test_access_token must be set when gateway auth is required" ) response = httpx.get( diff --git a/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py b/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py index da4c4a5cb..b30b30083 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py @@ -10,7 +10,6 @@ from http import HTTPStatus import pytest - from policyengine_api_simulation_client.models import ( BudgetWindowBatchSubmitResponse, BudgetWindowResult, diff --git a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py index b636f165b..942539177 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py @@ -12,7 +12,6 @@ from typing import Protocol, TypeGuard, cast import pytest - from policyengine_api_simulation_client import AuthenticatedClient, Client from policyengine_api_simulation_client.api.default import ( get_job_status_jobs_job_id_get, diff --git a/projects/policyengine-apis-integ/tests/simulation/test_ping.py b/projects/policyengine-apis-integ/tests/simulation/test_ping.py index ce990c287..71e9ae0b7 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_ping.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_ping.py @@ -1,9 +1,9 @@ -from policyengine_api_simulation_client import Client, AuthenticatedClient -from policyengine_api_simulation_client.api.default import ping_ping_post -from policyengine_api_simulation_client.models import PingRequest, PingResponse -from policyengine_api_simulation_client.errors import UnexpectedStatus import backoff import httpx +from policyengine_api_simulation_client import AuthenticatedClient, Client +from policyengine_api_simulation_client.api.default import ping_ping_post +from policyengine_api_simulation_client.errors import UnexpectedStatus +from policyengine_api_simulation_client.models import PingRequest, PingResponse @backoff.on_exception( diff --git a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py index 7e37009ac..6a2c19a21 100644 --- a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py +++ b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py @@ -3,15 +3,14 @@ import json import pytest - from policyengine_api_simulation_client import models as generated_models from policyengine_api_simulation_client.models import ( BudgetWindowAnnualImpact, + SimulationRequest, + SingleYearMacroOutput, SPMCapability, SPMComparisonProvenance, SPMSelection, - SingleYearMacroOutput, - SimulationRequest, ) diff --git a/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py b/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py index e2c8accd3..183b50d9f 100644 --- a/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py +++ b/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py @@ -7,7 +7,6 @@ import httpx import pytest - PROTECTED_PATH = "/simulate/economy/comparison" PROBE_PAYLOAD = { "country": "__authentication_probe__", 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 a3f332b99..0791fdf30 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py @@ -2,18 +2,17 @@ from __future__ import annotations -from functools import lru_cache import logging +from functools import lru_cache +import jwt from fastapi import Depends, HTTPException, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer -import jwt from policyengine_observability import ObservabilityRuntime from policyengine_simulation_entry.config import Settings from policyengine_simulation_entry.schemas import CallerIdentity - _bearer = HTTPBearer(auto_error=False) logger = logging.getLogger(__name__) 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 c2eef857e..f909e8455 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 @@ -367,15 +367,7 @@ def _prepare_automatic_report( ) evaluation_id = uuid5( AUTOMATIC_REPORT_NAMESPACE, - "|".join( - ( - self._settings.environment, - "economy", - production_identity, - version, - manifest_sha256, - ) - ), + f"{self._settings.environment}|economy|{production_identity}|{version}|{manifest_sha256}", ) adapted = adapt_annual_comparison( request_payload, diff --git a/projects/policyengine-simulation-entry/tests/conftest.py b/projects/policyengine-simulation-entry/tests/conftest.py index 524b7142b..560c1d6ef 100644 --- a/projects/policyengine-simulation-entry/tests/conftest.py +++ b/projects/policyengine-simulation-entry/tests/conftest.py @@ -5,7 +5,6 @@ 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 diff --git a/projects/policyengine-simulation-entry/tests/test_app.py b/projects/policyengine-simulation-entry/tests/test_app.py index 4c3a43ee2..d50e4e7f9 100644 --- a/projects/policyengine-simulation-entry/tests/test_app.py +++ b/projects/policyengine-simulation-entry/tests/test_app.py @@ -15,14 +15,6 @@ ComparisonRunLifecycleStatus, SimulationRole, ) -from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER -from stage12_fixtures import ( - EVALUATION_ID, - comparison_report, - comparison_simulation, - eligible_payload, -) - from policyengine_simulation_entry import app as app_module from policyengine_simulation_entry.app import create_app from policyengine_simulation_entry.backend import ( @@ -34,6 +26,13 @@ TemporaryStage12DispatchFailed, TemporaryStage12UnsupportedRequest, ) +from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER +from stage12_fixtures import ( + EVALUATION_ID, + comparison_report, + comparison_simulation, + eligible_payload, +) def response( diff --git a/projects/policyengine-simulation-entry/tests/test_auth.py b/projects/policyengine-simulation-entry/tests/test_auth.py index 1a1a552b9..acd697e40 100644 --- a/projects/policyengine-simulation-entry/tests/test_auth.py +++ b/projects/policyengine-simulation-entry/tests/test_auth.py @@ -6,16 +6,14 @@ import jwt import pytest +from conftest import FakeBackend, make_settings from cryptography.hazmat.primitives.asymmetric import rsa from fastapi.security import HTTPAuthorizationCredentials from fastapi.testclient import TestClient - from policyengine_simulation_entry import auth as auth_module from policyengine_simulation_entry.app import create_app from policyengine_simulation_entry.schemas import CallerIdentity -from conftest import FakeBackend, make_settings - @pytest.fixture def signing_keys(): diff --git a/projects/policyengine-simulation-entry/tests/test_backend.py b/projects/policyengine-simulation-entry/tests/test_backend.py index 6c0ceee99..85094b6c0 100644 --- a/projects/policyengine-simulation-entry/tests/test_backend.py +++ b/projects/policyengine-simulation-entry/tests/test_backend.py @@ -4,16 +4,14 @@ import httpx import pytest +from conftest import make_settings from policyengine_observability import REQUEST_ID_HEADER - from policyengine_simulation_entry.backend import ( BackendAuthenticationError, BackendUnavailable, OldGatewayBackend, ) -from conftest import make_settings - @pytest.mark.asyncio @pytest.mark.parametrize( diff --git a/projects/policyengine-simulation-entry/tests/test_cloud_run_revision_cleanup.py b/projects/policyengine-simulation-entry/tests/test_cloud_run_revision_cleanup.py index 4790de206..cd6bfd899 100644 --- a/projects/policyengine-simulation-entry/tests/test_cloud_run_revision_cleanup.py +++ b/projects/policyengine-simulation-entry/tests/test_cloud_run_revision_cleanup.py @@ -7,7 +7,6 @@ import textwrap from pathlib import Path - REPOSITORY_ROOT = Path(__file__).resolve().parents[3] CLEANUP_SCRIPT = ( REPOSITORY_ROOT diff --git a/projects/policyengine-simulation-entry/tests/test_config.py b/projects/policyengine-simulation-entry/tests/test_config.py index 127d9dff1..dc4252119 100644 --- a/projects/policyengine-simulation-entry/tests/test_config.py +++ b/projects/policyengine-simulation-entry/tests/test_config.py @@ -2,7 +2,6 @@ import pytest from conftest import make_settings - from policyengine_simulation_entry.config import ConfigurationError 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 14373c3b4..1c85eba93 100644 --- a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py +++ b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py @@ -12,7 +12,6 @@ from pathlib import Path import pytest - from workflow_script_support import ( WORKFLOWS_DIR, FakeCommands, diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 2a86571f1..9c0ae2556 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -1,7 +1,7 @@ from __future__ import annotations -from copy import deepcopy import json +from copy import deepcopy from pathlib import Path from typing import Literal, TypedDict, cast diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index a8b2cce35..c1b8ef4a9 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -6,12 +6,11 @@ import pytest from policyengine_simulation_contract.stage12_execution import SimulationRole -from stage12_fixtures import eligible_payload, worker - from policyengine_simulation_entry.stage12_adapter import ( ComparisonSkipReason, adapt_annual_comparison, ) +from stage12_fixtures import eligible_payload, worker EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_backend.py b/projects/policyengine-simulation-entry/tests/test_stage12_backend.py index 89b76c20c..55cdd73fe 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_backend.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_backend.py @@ -13,8 +13,6 @@ ComparisonReportRecord, ComparisonRunLifecycleStatus, ) -from stage12_fixtures import eligible_payload, worker - from policyengine_simulation_entry import stage12_backend as backend_module from policyengine_simulation_entry.stage12_backend import ( ModalReportInvoker, @@ -22,6 +20,7 @@ TemporaryStage12DispatchFailed, TemporaryStage12UnsupportedRequest, ) +from stage12_fixtures import eligible_payload, worker OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000012" diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py b/projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py index abc4a3ea1..b3e84d470 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py @@ -2,8 +2,8 @@ from __future__ import annotations -from pathlib import Path import re +from pathlib import Path REPOSITORY_ROOT = Path(__file__).resolve().parents[3] RUNTIME_SOURCE_ROOTS = ( diff --git a/projects/policyengine-simulation-entry/tests/test_workflow_scripts.py b/projects/policyengine-simulation-entry/tests/test_workflow_scripts.py index 4a1544da6..ab46e34cb 100644 --- a/projects/policyengine-simulation-entry/tests/test_workflow_scripts.py +++ b/projects/policyengine-simulation-entry/tests/test_workflow_scripts.py @@ -8,7 +8,6 @@ from pathlib import Path import pytest - from workflow_script_support import ( REPOSITORY_ROOT, WORKFLOWS_DIR, diff --git a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py index d260dd5ad..0a6c1ca5a 100644 --- a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py +++ b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py @@ -25,13 +25,11 @@ def install_identity_stubs(monkeypatch): to the installed capability resolver for tests that intentionally exercise the installed release. """ - from policyengine.provenance import manifest as manifest_module - # This module imports ``get_release_manifest`` by value while constructing # country models. Import it before replacing the manifest module's function # so full-suite ordering cannot bind this test's stub into production code. import policyengine.tax_benefit_models.common.model_version # noqa: F401 - + from policyengine.provenance import manifest as manifest_module from policyengine_simulation_executor import release_bundle bundle = SimpleNamespace( diff --git a/projects/policyengine-simulation-executor/scripts/dump_package_version.py b/projects/policyengine-simulation-executor/scripts/dump_package_version.py index 288fed6b8..afa2d0403 100644 --- a/projects/policyengine-simulation-executor/scripts/dump_package_version.py +++ b/projects/policyengine-simulation-executor/scripts/dump_package_version.py @@ -1,6 +1,7 @@ -import tomllib import sys +import tomllib + def get_package_version_from_lock(package_name, lock_file_path="poetry.lock"): try: diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index 4e96b2276..4e5786b01 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -10,6 +10,14 @@ import os from pathlib import Path +from policyengine_simulation_executor.hf_access_validation import ( + HF_MODAL_SECRET_NAME, + HF_RUNTIME_ENV_NAME, +) +from policyengine_simulation_executor.release_bundle import ( + get_bundled_country_model_version, + get_bundled_package_version, +) from policyengine_simulation_observability.observability import ( init_process_observability, modal_image_environment, @@ -23,21 +31,13 @@ from policyengine_simulation_observability.telemetry import ( normalize_observability_context, ) - -import modal -from policyengine_simulation_executor.hf_access_validation import ( - HF_MODAL_SECRET_NAME, - HF_RUNTIME_ENV_NAME, -) -from policyengine_simulation_executor.release_bundle import ( - get_bundled_country_model_version, - get_bundled_package_version, -) from src.modal._image_setup import fetch_artifacts, snapshot_models from src.modal.bundle_data import bundle_data_install_command from src.modal.logging_redaction import redact_params_for_logging from src.modal.static_runtime_files import add_static_runtime_files +import modal + def _version_from_env_or_local_bundle_package(env_var: str, package: str) -> str: value = os.environ.get(env_var) 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 3707f5dfe..5403375b0 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py @@ -4,12 +4,12 @@ from typing import Any -import modal 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 -from policyengine_simulation_observability.stages import BUDGET_WINDOW_STAGES, Stage + +import modal def run_budget_window_batch_impl( diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_context.py b/projects/policyengine-simulation-executor/src/modal/budget_window_context.py index a0a37bba4..ce91f8073 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_context.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_context.py @@ -9,7 +9,6 @@ BudgetWindowBatchRequest, PolicyEngineBundle, ) - from src.modal.fanout import build_child_payload BATCH_ONLY_FIELDS = { diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py index 8eabb3801..764e38674 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py @@ -5,13 +5,12 @@ from decimal import Decimal from typing import Any -from policyengine_simulation_contract.spm import validate_spm_result - from policyengine_simulation_contract.gateway_models import ( BudgetWindowAnnualImpact, BudgetWindowResult, BudgetWindowTotals, ) +from policyengine_simulation_contract.spm import validate_spm_result # The UK microsimulation has no state/province fiscal layer, so worker child # results for ``country="uk"`` never emit ``state_tax_revenue_impact``. The @@ -25,7 +24,7 @@ OPTIONAL_BUDGET_KEYS = ("state_tax_revenue_impact",) -def _as_decimal(value: float | int) -> Decimal: +def _as_decimal(value: float) -> Decimal: """Convert an annual impact float to Decimal without reintroducing binary-float quantisation noise. ``Decimal(str(...))`` is the canonical idiom because it serialises the float to its shortest round-trippable 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 1aa7bb7a2..602083f1c 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py @@ -22,8 +22,6 @@ from policyengine_simulation_contract.spm import spm_error_detail from policyengine_simulation_observability.errors import log_and_redact_exception from policyengine_simulation_observability.stages import BUDGET_WINDOW_STAGES, Stage - -import modal from src.modal.budget_window_context import ( BudgetWindowBatchContext, ChildSimulationHandle, @@ -34,6 +32,8 @@ extract_annual_impact, ) +import modal + # 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/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py index db2d49aff..34d3b0833 100644 --- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py @@ -17,14 +17,14 @@ from __future__ import annotations -import modal - from src.modal.app import ( build_runtime_simulation_image, gcp_secret, hf_secret, ) +import modal + app = modal.App("policyengine-simulation-precompute") precompute_image = build_runtime_simulation_image().add_local_python_source( diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 9e5ed5cc2..83328b08b 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -30,21 +30,21 @@ combine_spm_results, spm_error_detail, ) +from policyengine_simulation_executor.national_partition import ( + national_region_groups, +) +from policyengine_simulation_executor.segmented_national_reduce import ( + build_national_output, +) from policyengine_simulation_observability.errors import log_and_redact_exception from policyengine_simulation_observability.stages import ( SEGMENTED_NATIONAL_STAGES, Stage, ) from policyengine_simulation_observability.telemetry import split_internal_payload +from src.modal.fanout import build_child_payload, next_backoff import modal -from policyengine_simulation_executor.national_partition import ( - national_region_groups, -) -from policyengine_simulation_executor.segmented_national_reduce import ( - build_national_output, -) -from src.modal.fanout import build_child_payload, next_backoff logger = logging.getLogger(__name__) diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index 10b7673d0..4318b5210 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -20,16 +20,16 @@ from pathlib import Path -import modal -from src.modal.app import build_runtime_simulation_image -from src.modal.static_runtime_files import add_static_runtime_files - from policyengine_simulation_executor.release_bundle import ( resolve_local_bundle_dataset_path, ) from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata_from_hdf, ) +from src.modal.app import build_runtime_simulation_image +from src.modal.static_runtime_files import add_static_runtime_files + +import modal app = modal.App("policyengine-simulation-executor-smoke") @@ -136,7 +136,6 @@ def smoke_import_executor() -> dict: # smoke exists to catch. import policyengine_simulation_contract import policyengine_simulation_observability - from policyengine_simulation_executor.simulation_runtime import ( # noqa: F401 run_simulation_impl, ) 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 a3a8a0350..f7715d5e5 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py @@ -10,8 +10,8 @@ import argparse import os +from collections.abc import Mapping from datetime import datetime, timezone -from typing import Mapping from policyengine_simulation_executor.artifact_store import ArtifactStore from policyengine_simulation_executor.precompute_models import DeployedMarker diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py index 0d267bb90..4717bb628 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py @@ -7,7 +7,6 @@ from pathlib import Path import modal - from policyengine_simulation_contract.stage12_manifest import ( V1_ROUTING_STATE_NAME, V2_VERSION_MANIFEST_NAME, diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py index 649a652a6..3462a4e8a 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py @@ -1,11 +1,12 @@ """Publish the Modal routing state after a simulation API deployment.""" import argparse +from collections.abc import Iterable from copy import deepcopy +from typing import NotRequired, TypedDict import modal from packaging.version import InvalidVersion, Version -from typing import Iterable, TypedDict, NotRequired POLICYENGINE_VERSION_DICT_NAME = "simulation-api-policyengine-versions" US_VERSION_DICT_NAME = "simulation-api-us-versions" 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 348076fd1..fc088bd2e 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 @@ -3,14 +3,13 @@ from __future__ import annotations import argparse -from datetime import datetime, timezone import json +from datetime import datetime, timezone from pathlib import Path from typing import Any from uuid import uuid4 import modal - from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_manifest import ( ManifestText, diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 357670a2a..076f9dce0 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -13,6 +13,15 @@ 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, + HF_RUNTIME_ENV_NAME, +) +from policyengine_simulation_executor.stage12_bundle import ( + assert_expected_bundle_values, + assertion_values, + load_stage12_bundle, +) from policyengine_simulation_observability.observability import ( init_process_observability, modal_image_environment, @@ -26,20 +35,11 @@ from policyengine_simulation_observability.telemetry import ( normalize_observability_context, ) - -import modal -from policyengine_simulation_executor.hf_access_validation import ( - HF_MODAL_SECRET_NAME, - HF_RUNTIME_ENV_NAME, -) -from policyengine_simulation_executor.stage12_bundle import ( - assert_expected_bundle_values, - assertion_values, - load_stage12_bundle, -) from src.modal.bundle_data import bundle_data_install_command from src.modal.static_runtime_files import add_static_runtime_files +import modal + 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/modal/v2_static_runtime_files_smoke.py b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py index fa9166dd3..4e5f5ac9a 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py @@ -1,9 +1,6 @@ # pyright: reportAttributeAccessIssue=false, reportFunctionMemberAccess=false """Pre-merge smoke for static files inside the exact UK Stage 12 image.""" -import modal -from src.modal.v2_app import uk_worker_image - from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, ) @@ -14,6 +11,9 @@ detect_uk_local_authority_metadata_from_hdf, load_uk_local_authority_resources, ) +from src.modal.v2_app import uk_worker_image + +import modal app = modal.App("policyengine-stage12-v2-static-runtime-files-smoke") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py index cb8a2744f..98f1a49b6 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py @@ -1,5 +1,6 @@ from fastapi import FastAPI from policyengine_observability import ObservabilityRuntime + from .simulation import create_router """ diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py index 5cdcffba9..9eafa91ab 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py @@ -36,7 +36,7 @@ import json from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Optional +from typing import Any DATASET_KEY_SCHEMA = "ds1" BASELINE_KEY_SCHEMA = "bl1" @@ -70,11 +70,11 @@ def dataset_key( year: int, data_version: str, data_artifact_revision: str, - source_sha256: Optional[str], - data_build_fingerprint: Optional[str], + source_sha256: str | None, + data_build_fingerprint: str | None, model_version: str, policyengine_version: str, - spm: Optional[dict] = None, + spm: dict | None = None, ) -> str: """Digest identifying one single-year dataset artifact. @@ -106,11 +106,11 @@ def baseline_key( *, country: str, region: str, - scope_key: Optional[str], + scope_key: str | None, dataset_digest: str, model_version: str, policyengine_version: str, - spm: Optional[dict] = None, + spm: dict | None = None, policy: str = CURRENT_LAW_POLICY, ) -> str: """Digest identifying one precomputed baseline simulation artifact. @@ -206,11 +206,11 @@ class DatasetArtifactIdentity: year: int data_version: str data_artifact_revision: str - source_sha256: Optional[str] - data_build_fingerprint: Optional[str] + source_sha256: str | None + data_build_fingerprint: str | None model_version: str policyengine_version: str - spm: Optional[dict] = None + spm: dict | None = None @property def digest(self) -> str: @@ -242,9 +242,9 @@ class BaselineArtifactIdentity: country: str region: str - scope_key: Optional[str] + scope_key: str | None dataset: DatasetArtifactIdentity - spm: Optional[dict] = None + spm: dict | None = None @property def digest(self) -> str: @@ -273,7 +273,7 @@ def store_path(self) -> str: return baseline_artifact_path(self.country, self.digest, self.storage_id) -def _receipt_source_sha256(country: str, data_version: str) -> Optional[str]: +def _receipt_source_sha256(country: str, data_version: str) -> str | None: """Content hash of the installed multi-year source, when certifiable. Only trusted when the receipt entry's version matches the bundle's data @@ -294,7 +294,7 @@ def _receipt_source_sha256(country: str, data_version: str) -> Optional[str]: return None -def _certification_fingerprint(country: str) -> Optional[str]: +def _certification_fingerprint(country: str) -> str | None: from policyengine.provenance.manifest import get_release_manifest certification = getattr(get_release_manifest(country), "certification", None) @@ -347,8 +347,8 @@ def collect_baseline_identity( year: int, *, region: str, - scope_key: Optional[str], - spm: Optional[dict] = None, + scope_key: str | None, + spm: dict | None = None, ) -> BaselineArtifactIdentity: from policyengine_simulation_executor.spm import normalize_runtime_spm diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py index 1f84b91de..b16b155f8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py @@ -24,7 +24,7 @@ import os from collections.abc import Mapping from pathlib import Path -from typing import Any, Optional +from typing import Any from google.api_core.exceptions import NotFound, PreconditionFailed @@ -33,7 +33,7 @@ ARTIFACT_BUCKET_ENV = "POLICYENGINE_ARTIFACT_BUCKET" -def resolve_bucket_name(explicit: Optional[str] = None) -> str: +def resolve_bucket_name(explicit: str | None = None) -> str: bucket = explicit or os.environ.get(ARTIFACT_BUCKET_ENV) if not bucket: raise RuntimeError( @@ -56,7 +56,7 @@ def _default_client(): class ArtifactStore: - def __init__(self, bucket_name: Optional[str] = None, *, client: Any = None): + def __init__(self, bucket_name: str | None = None, *, client: Any = None): self.bucket_name = resolve_bucket_name(bucket_name) self._client = client @@ -137,7 +137,7 @@ def _upload_json(self, path: str, payload: Mapping[str, Any], *, overwrite: bool return False return True - def read_json(self, path: str) -> Optional[Mapping[str, Any]]: + def read_json(self, path: str) -> Mapping[str, Any] | None: try: payload = json.loads(self._blob(path).download_as_bytes()) except NotFound: @@ -152,7 +152,7 @@ def write_manifest(self, payload: Mapping[str, Any]) -> str: self._upload_json(artifact_keys.manifest_path(digest), payload, overwrite=False) return digest - def read_manifest(self, digest: str) -> Optional[Mapping[str, Any]]: + def read_manifest(self, digest: str) -> Mapping[str, Any] | None: return self.read_json(artifact_keys.manifest_path(digest)) # --- Deployed-environment markers (last-writer-wins) ------------------ @@ -164,5 +164,5 @@ def write_deployed_marker( artifact_keys.deployed_marker_path(environment), payload, overwrite=True ) - def read_deployed_marker(self, environment: str) -> Optional[Mapping[str, Any]]: + def read_deployed_marker(self, environment: str) -> Mapping[str, Any] | None: return self.read_json(artifact_keys.deployed_marker_path(environment)) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py index a2af47817..e66a59092 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py @@ -21,11 +21,10 @@ from __future__ import annotations import logging -from typing import Any, Optional - -from pydantic import PrivateAttr +from typing import Any from policyengine.core import Simulation +from pydantic import PrivateAttr from policyengine_simulation_executor import artifact_keys @@ -47,10 +46,10 @@ def qualifying_baseline_identity( country: str, dataset_is_default: bool, policy: Any, - region_code: Optional[str], + region_code: str | None, scoping_strategy: Any, year: int, -) -> Optional[artifact_keys.BaselineArtifactIdentity]: +) -> artifact_keys.BaselineArtifactIdentity | None: """Full artifact identity for a qualifying baseline, else None. Qualifying means the simulation's bytes are a pure function of the @@ -105,10 +104,10 @@ def deterministic_baseline_id( country: str, dataset_is_default: bool, policy: Any, - region_code: Optional[str], + region_code: str | None, scoping_strategy: Any, year: int, -) -> Optional[str]: +) -> str | None: """The runtime wrapper: id for a qualifying baseline, else None. Unlike the writer, the read path must never fail a request over @@ -144,10 +143,10 @@ class ArtifactBaselineSimulation(Simulation): """A Simulation whose ``ensure()`` validates what it loads.""" _computed_this_process: bool = PrivateAttr(default=False) - _artifact_outcome: Optional[str] = PrivateAttr(default=None) + _artifact_outcome: str | None = PrivateAttr(default=None) @property - def artifact_outcome(self) -> Optional[str]: + def artifact_outcome(self) -> str | None: """Set by ``ensure()``; None until then.""" return self._artifact_outcome 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 f50ea3b0b..beaca988d 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 @@ -2,10 +2,10 @@ from __future__ import annotations -from typing import Any, Optional +from typing import Any -from pydantic import BaseModel, ConfigDict from policyengine_simulation_contract.spm import SPMSelection +from pydantic import BaseModel, ConfigDict from policyengine_simulation_executor.simulation_macro_output import ( SingleYearMacroOutput, @@ -16,20 +16,20 @@ class SimulationOptions(BaseModel): """Legacy request schema name kept for generated clients.""" country: str - spm: Optional[SPMSelection] = None - scope: Optional[str] = None - time_period: Optional[str | int] = None - reform: Optional[dict[str, Any]] = None - baseline: Optional[dict[str, Any]] = None - region: Optional[str] = None - region_group: Optional[list[str]] = None - title: Optional[str] = None - include_cliffs: Optional[bool] = None - model_version: Optional[str] = None + spm: SPMSelection | None = None + scope: str | None = None + time_period: str | int | None = None + reform: dict[str, Any] | None = None + baseline: dict[str, Any] | None = None + region: str | None = None + region_group: list[str] | None = None + title: str | None = None + include_cliffs: bool | None = None + model_version: str | None = None # Accepted for parity with the gateway contract. The synchronous # surface always runs monolithically, so segmented=false is trivially # honored and segmented=true has no effect here. - segmented: Optional[bool] = None + segmented: bool | None = None model_config = ConfigDict(extra="forbid") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/generate_openapi.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/generate_openapi.py index 3c8047638..5d43c0925 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/generate_openapi.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/generate_openapi.py @@ -1,7 +1,8 @@ -from .main import app import json from pathlib import Path +from .main import app + # Write OpenAPI spec directly to file output_path = Path(__file__).parent.parent.parent / "artifacts" / "openapi.json" output_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/hf_access_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/hf_access_validation.py index 51c4e4919..2565dc13e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/hf_access_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/hf_access_validation.py @@ -8,9 +8,9 @@ from typing import Literal from huggingface_hub import get_hf_file_metadata, hf_hub_url +from policyengine_simulation_contract.hf_dataset import parse_hf_dataset_uri from pydantic import BaseModel, ConfigDict, Field -from policyengine_simulation_contract.hf_dataset import parse_hf_dataset_uri from policyengine_simulation_executor.release_bundle import ( get_country_release_bundle, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py index 4ede15ff2..b0b0d7383 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py @@ -1,16 +1,18 @@ +import logging from contextlib import asynccontextmanager + from fastapi import FastAPI -from policyengine_fastapi.exit import exit -from policyengine_simulation_executor import initialize -from policyengine_simulation_observability.observability import ( - init_simulation_observability, -) from policyengine_fastapi import ping +from policyengine_fastapi.exit import exit from policyengine_fastapi.health import ( HealthRegistry, HealthSystemReporter, ) -import logging +from policyengine_simulation_observability.observability import ( + init_simulation_observability, +) + +from policyengine_simulation_executor import initialize """ specific example instantiation of the app configured by a .env file 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 9a2f5c511..e58707a46 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -34,7 +34,7 @@ import logging import time -from typing import Callable +from collections.abc import Callable from policyengine_simulation_executor.artifact_keys import ( BaselineArtifactIdentity, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py index fd5332fd2..a6abc66e2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py @@ -20,7 +20,7 @@ from __future__ import annotations -from typing import Any, Literal, Optional, Protocol +from typing import Any, Literal, Protocol from pydantic import BaseModel, ConfigDict, Field @@ -148,7 +148,7 @@ class BaselineComputeResult(_StrictModel): year: int group: list[str] simulation_id: str - outcome: Optional[ArtifactOutcome] + outcome: ArtifactOutcome | None uploaded: bool compute_seconds: float size_bytes: int diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py index 46a8e455d..71fcfbfdf 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py @@ -25,12 +25,10 @@ import logging from typing import Any -from policyengine_observability import ObservabilityRuntime - import pandas as pd from microdf import MicroDataFrame - from policyengine.core import Simulation +from policyengine_observability import ObservabilityRuntime from policyengine_simulation_executor.simulation_microdata import ( rebuild_entity_frame, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py index 09c3407d9..fd76d9751 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py @@ -1,7 +1,8 @@ from enum import Enum from functools import lru_cache -from pydantic_settings import BaseSettings, SettingsConfigDict + from pydantic import field_validator +from pydantic_settings import BaseSettings, SettingsConfigDict class Environment(Enum): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py index 884b4b3fe..78d8b731d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py @@ -5,11 +5,11 @@ from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.spm import spm_error_detail -from policyengine_simulation_executor.simulation_runtime import run_simulation_impl from policyengine_simulation_executor.compat_models import ( EconomyComparison, SimulationOptions, ) +from policyengine_simulation_executor.simulation_runtime import run_simulation_impl logger = logging.getLogger(__file__) 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 5805c8bee..cd0636b7f 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 @@ -8,18 +8,20 @@ from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata - -from policyengine_simulation_executor import simulation_output_budget -from policyengine_simulation_executor import simulation_output_cliff -from policyengine_simulation_executor import simulation_output_distribution -from policyengine_simulation_executor import simulation_output_geographic -from policyengine_simulation_executor import simulation_output_inequality -from policyengine_simulation_executor import simulation_output_labor -from policyengine_simulation_executor import simulation_output_poverty from policyengine_simulation_observability.stages import ( ANNUAL_IMPACT_STAGES, Stage, ) + +from policyengine_simulation_executor import ( + simulation_output_budget, + simulation_output_cliff, + simulation_output_distribution, + simulation_output_geographic, + simulation_output_inequality, + simulation_output_labor, + simulation_output_poverty, +) from policyengine_simulation_executor.release_bundle import get_country_release_bundle from policyengine_simulation_executor.simulation_macro_output import ( BudgetaryImpact, 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 21dff1d00..3e71d8df5 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 @@ -34,6 +34,7 @@ def _required_uk_geography_lookup_csv_path(spec: UKGeographyAssetSpec) -> str: from policyengine.outputs.uk_geography_impact import ( resolve_uk_geography_lookup_csv_path, ) + from policyengine_simulation_executor.simulation_runtime import ( setup_gcp_credentials, ) @@ -180,12 +181,8 @@ def _should_build_us_congressional_district_impact(region_code: str | None) -> b if region_code is None: return True normalized_region_code = region_code.lower() - return ( - normalized_region_code == "us" - or normalized_region_code.startswith("state/") - # A region group is a union of whole states; the CD table computed over - # its households is exactly its member states' districts. - or normalized_region_code.startswith("region_group/") + return normalized_region_code == "us" or normalized_region_code.startswith( + ("state/", "region_group/") ) @@ -275,6 +272,7 @@ def build_uk_local_authority_impact( from policyengine.outputs.uk_geography_impact import ( compute_longwise_uk_geography_impacts, ) + from policyengine_simulation_executor.uk_local_authority_metadata import ( load_uk_local_authority_resources, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index 5018b0f79..6004f9283 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -13,18 +13,20 @@ import os import tempfile import time +from collections.abc import Iterator from dataclasses import dataclass from datetime import datetime from importlib import import_module -from typing import Any, Iterator +from typing import Any from policyengine_observability import ObservabilityRuntime - from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri from policyengine_simulation_observability.stages import ( ANNUAL_IMPACT_STAGES, Stage, ) +from policyengine_simulation_observability.telemetry import split_internal_payload + from policyengine_simulation_executor.release_bundle import ( get_country_release_bundle, resolve_bundle_dataset_name, @@ -32,7 +34,6 @@ from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, ) -from policyengine_simulation_observability.telemetry import split_internal_payload logger = logging.getLogger(__name__) @@ -689,7 +690,6 @@ def _build_simulation( ArtifactBaselineSimulation, deterministic_baseline_id, ) - from policyengine_simulation_executor.spm import normalize_runtime_spm selection = normalize_runtime_spm(params) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index 45daa3a0c..d76ad3fcb 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -10,11 +10,11 @@ SPMProvenance, SPMRuntimeVersions, SPMSelection, - validate_spm_calculation_provenance, build_spm_comparison_provenance, build_spm_provenance, resolve_spm_selection, spm_error_detail, + validate_spm_calculation_provenance, ) 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 ff24ea350..5b07ffa1c 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 @@ -13,6 +13,7 @@ import pandas as pd import pyarrow as pa import pyarrow.parquet as pq +from policyengine_simulation_contract.spm import SPMProvenance from policyengine_simulation_contract.stage12_execution import ( SIMULATION_PARQUET_PAYLOAD_CONTRACT, AggregateReportArtifactPayload, @@ -24,7 +25,6 @@ SimulationArtifactDescriptor, stage12_output_plan_sha256, ) -from policyengine_simulation_contract.spm import SPMProvenance from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_executor.artifact_store import ArtifactStore diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py index 10d9cfe35..0a7e4984d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py @@ -2,9 +2,9 @@ from __future__ import annotations +import math from collections.abc import Mapping, Sequence from dataclasses import dataclass -import math import numpy as np import pandas as pd diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py index 473d1ea82..49c473f36 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py @@ -16,8 +16,6 @@ OTelConfig, ServiceIdentity, ) -from pydantic import BaseModel, ConfigDict - from policyengine_simulation_contract.stage12_execution import ( ArtifactMediaType, ArtifactReference, @@ -26,6 +24,8 @@ SimulationArtifactDescriptor, stage12_output_plan_sha256, ) +from pydantic import BaseModel, ConfigDict + from policyengine_simulation_executor.simulation_microdata import ( rebuild_entity_frame, ) 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 99055a349..908141f8d 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 @@ -6,14 +6,14 @@ from contextlib import nullcontext from datetime import UTC, datetime from typing import Any -from policyengine_observability import ObservabilityRuntime -from policyengine_simulation_observability.stages import Stage, StagePlan +from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( ComparisonReportRecord, ResultComparisonStatus, Stage12InvocationContext, ) +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 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 cc483437f..88a89a987 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 @@ -11,6 +11,7 @@ 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, @@ -21,7 +22,6 @@ Stage12InvocationContext, stage12_output_plan_sha256, ) -from policyengine_simulation_contract.spm import SPMProvenance from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_observability.stages import ( STAGE12_SIMULATION_STAGES, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py index 5f5114237..8f97f2237 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py @@ -10,6 +10,7 @@ from functools import lru_cache from importlib.resources import files from types import MappingProxyType + import pandas as pd from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, diff --git a/projects/policyengine-simulation-executor/tests/test_artifact_keys.py b/projects/policyengine-simulation-executor/tests/test_artifact_keys.py index ef648b74f..c7b1dcffb 100644 --- a/projects/policyengine-simulation-executor/tests/test_artifact_keys.py +++ b/projects/policyengine-simulation-executor/tests/test_artifact_keys.py @@ -12,7 +12,6 @@ from types import SimpleNamespace import pytest - from fixtures.identity_stubs import install_identity_stubs from fixtures.wrapper_spm import ( wrapper_storage_id as _wrapper_storage_id, diff --git a/projects/policyengine-simulation-executor/tests/test_artifact_store.py b/projects/policyengine-simulation-executor/tests/test_artifact_store.py index 0310277df..d9329aae6 100644 --- a/projects/policyengine-simulation-executor/tests/test_artifact_store.py +++ b/projects/policyengine-simulation-executor/tests/test_artifact_store.py @@ -10,7 +10,6 @@ import pytest from google.api_core.exceptions import NotFound, PreconditionFailed - from policyengine_simulation_executor import artifact_keys as ak from policyengine_simulation_executor.artifact_store import ( ARTIFACT_BUCKET_ENV, diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index 6f9fdfaf2..8c2f06865 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -13,10 +13,9 @@ import pandas as pd import pytest from policyengine.core import Simulation -from pydantic import PrivateAttr - from policyengine_simulation_executor import artifact_keys as ak from policyengine_simulation_executor import baseline_artifacts as ba +from pydantic import PrivateAttr def _region_group(): @@ -351,7 +350,6 @@ class TestBuildSimulationSelection: @pytest.fixture def wired(self, monkeypatch): from policyengine import core as policyengine_core - from policyengine_simulation_executor import simulation_runtime as sr recorded = SimpleNamespace( @@ -456,7 +454,6 @@ def test_non_qualifying_uses_plain_simulation(self, wired): def test_bounded_reform_build_keeps_current_law_artifact_reuse(monkeypatch): from policyengine import core as policyengine_core - from policyengine_simulation_executor import simulation_runtime as sr from policyengine_simulation_executor import spm as executor_spm diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py index 34960eeab..2e5b26e40 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py @@ -9,7 +9,6 @@ from pathlib import Path import pytest - from native_spm_support import materialize_native_household pytestmark = pytest.mark.skipif( diff --git a/projects/policyengine-simulation-executor/tests/test_budget_reform_contract_route.py b/projects/policyengine-simulation-executor/tests/test_budget_reform_contract_route.py index e0724e85d..9b3a1005d 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_reform_contract_route.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_reform_contract_route.py @@ -22,7 +22,6 @@ from policyengine.utils.parametric_reforms import ( # noqa: E402 reform_dict_from_parameter_values, ) - from policyengine_simulation_executor.simulation_runtime import ( # noqa: E402 _compile_interval_policy, _country_module, diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py index d9f66db47..e0642c76e 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py @@ -6,13 +6,12 @@ import policyengine_simulation_contract.budget_window_state as state_module import pytest +import src.modal.budget_window_batch as batch_module +import src.modal.budget_window_scheduler as scheduler_module from policyengine_simulation_contract.gateway_models import ( BudgetWindowBatchRequest, PolicyEngineBundle, ) - -import src.modal.budget_window_batch as batch_module -import src.modal.budget_window_scheduler as scheduler_module from src.modal.budget_window_batch import run_budget_window_batch_impl @@ -28,7 +27,7 @@ def get(self, timeout: int = 0): event = self._events[0] if isinstance(event, TimeoutError): self._events.popleft() - raise TimeoutError() + raise TimeoutError self._events.popleft() if isinstance(event, Exception): diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_context.py b/projects/policyengine-simulation-executor/tests/test_budget_window_context.py index ccf337c33..1a1692dc0 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_context.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_context.py @@ -1,13 +1,13 @@ """Tests for budget-window batch context helpers.""" -from src.modal.budget_window_context import ( - build_batch_context, - build_child_simulation_request, -) from policyengine_simulation_contract.gateway_models import ( BudgetWindowBatchRequest, PolicyEngineBundle, ) +from src.modal.budget_window_context import ( + build_batch_context, + build_child_simulation_request, +) def _build_parent_payload(): diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_results.py b/projects/policyengine-simulation-executor/tests/test_budget_window_results.py index 5d2fca082..cd8319c56 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_results.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_results.py @@ -1,13 +1,12 @@ """Tests for budget-window batch result helpers.""" import pytest - +from policyengine_simulation_contract.gateway_models import BudgetWindowAnnualImpact from src.modal.budget_window_results import ( build_budget_window_result, extract_annual_impact, sum_annual_impacts, ) -from policyengine_simulation_contract.gateway_models import BudgetWindowAnnualImpact def test_extract_annual_impact_matches_v1_shape(): diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py index 1c0a0e032..ca41b7b8b 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py @@ -15,6 +15,8 @@ import policyengine_simulation_contract.budget_window_state as state_module import pytest +import src.modal.budget_window_batch as batch_module +import src.modal.budget_window_scheduler as scheduler_module from fastapi.testclient import TestClient from policyengine_simulation_contract.budget_window_state import ( BUDGET_WINDOW_JOB_DICT_NAME, @@ -25,9 +27,6 @@ from policyengine_simulation_gateway.testing import create_gateway_app from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER -import src.modal.budget_window_batch as batch_module -import src.modal.budget_window_scheduler as scheduler_module - SPM_SELECTION = SPMSelection( forecast_content_sha256="a" * 64, scenario="ce_trend", @@ -167,7 +166,7 @@ def get(self, timeout: int = 0): return self._result if self._polls == 0: self._polls += 1 - raise TimeoutError() + raise TimeoutError previous = self.runtime.current_parent_call_id self.runtime.current_parent_call_id = self.object_id diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index b091a9a23..d73c6828f 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -10,24 +10,23 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient - from policyengine_simulation_contract.gateway_models import ( - SimulationRequest, BudgetWindowBatchRequest, + SimulationRequest, ) from policyengine_simulation_contract.spm import ( SPMCapability, - SPMSelection, SPMInputError, - resolve_spm_selection, + SPMSelection, combine_spm_results, + resolve_spm_selection, ) from policyengine_simulation_executor import artifact_keys, simulation_runtime from policyengine_simulation_executor.simulation import create_router from policyengine_simulation_executor.spm import compact_spm_provenance from src.modal.budget_window_results import ( - extract_annual_impact, build_budget_window_result, + extract_annual_impact, ) from src.modal.fanout import build_child_payload @@ -602,8 +601,8 @@ def test_fanout_and_budget_window_preserve_settings_and_receipts(): def test_worker_build_passes_selection_to_baseline_and_reform(monkeypatch): import policyengine.core - from policyengine_simulation_executor import baseline_artifacts import policyengine_simulation_executor.spm as spm + from policyengine_simulation_executor import baseline_artifacts calls = [] monkeypatch.setattr(spm, "normalize_runtime_spm", lambda params: SELECTION) @@ -689,9 +688,9 @@ def test_gateway_poll_returns_structured_400(monkeypatch, code): def test_gateway_submission_uses_registry_capability_before_spawn(monkeypatch): + from policyengine_simulation_contract.gateway_models import PolicyEngineBundle from policyengine_simulation_gateway import endpoints from policyengine_simulation_gateway.testing import create_gateway_app - from policyengine_simulation_contract.gateway_models import PolicyEngineBundle route = SimpleNamespace( app_name="test-app", @@ -716,7 +715,7 @@ def test_gateway_submission_uses_registry_capability_before_spawn(monkeypatch): ) ), ) - monkeypatch.setattr(endpoints, "_job_metadata_store", lambda: {}) + monkeypatch.setattr(endpoints, "_job_metadata_store", dict) response = TestClient(create_gateway_app()).post( "/simulate/economy/comparison", json={"country": "us", "spm": {"geography_kind": "national"}}, @@ -758,10 +757,10 @@ def test_segmented_child_input_error_is_not_redacted_or_retried(monkeypatch): def test_budget_window_state_keeps_typed_failure_on_replay(): - from policyengine_simulation_contract.gateway_models import BudgetWindowBatchState from policyengine_simulation_contract.budget_window_state import ( build_batch_status_response, ) + from policyengine_simulation_contract.gateway_models import BudgetWindowBatchState from policyengine_simulation_gateway.responses import batch_status_response state = BudgetWindowBatchState( @@ -853,11 +852,11 @@ def test_optional_analysis_does_not_swallow_spm_input_errors(code): [({"as_of": None}, None), ({}, "2026-01-01")], ) def test_as_of_presence_survives_budget_parent(date_fields, expected_as_of): + from policyengine_simulation_contract.gateway_models import PolicyEngineBundle from policyengine_simulation_gateway.endpoints import ( - _resolve_request_spm, _build_budget_window_parent_payload, + _resolve_request_spm, ) - from policyengine_simulation_contract.gateway_models import PolicyEngineBundle request = BudgetWindowBatchRequest( country="us", diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py index 6b0e732e0..799197748 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py @@ -8,7 +8,6 @@ import os import pytest - from native_spm_support import materialize_native_household SOURCE = os.environ.get("SPM_NATIVE_SMOKE_SOURCE") @@ -118,11 +117,11 @@ def test_worker_baseline_reform_national_local_cache_and_receipts( def test_native_state_only_worker_requires_geography_and_explicit_national_works( native_dataset, ): + from policyengine_simulation_contract.spm import spm_error_detail from policyengine_simulation_executor.simulation_runtime import ( DatasetSelection, _build_simulation, ) - from policyengine_simulation_contract.spm import spm_error_detail native_dataset.data.household.drop(columns=["county_fips"], inplace=True) params = { @@ -154,8 +153,8 @@ def test_native_state_only_worker_requires_geography_and_explicit_national_works def test_actual_provider_unknown_metro_is_a_typed_input_error(): - from policyengine_simulation_executor.spm import normalize_runtime_spm from policyengine_simulation_contract.spm import SPMInputError + from policyengine_simulation_executor.spm import normalize_runtime_spm with pytest.raises(SPMInputError) as error: normalize_runtime_spm( diff --git a/projects/policyengine-simulation-executor/tests/test_gcp_credentials.py b/projects/policyengine-simulation-executor/tests/test_gcp_credentials.py index 0bd913446..530689baa 100644 --- a/projects/policyengine-simulation-executor/tests/test_gcp_credentials.py +++ b/projects/policyengine-simulation-executor/tests/test_gcp_credentials.py @@ -6,7 +6,6 @@ import os import pytest - from policyengine_simulation_executor.simulation_runtime import ( _normalize_credentials_blob, setup_gcp_credentials, diff --git a/projects/policyengine-simulation-executor/tests/test_hf_access_validation.py b/projects/policyengine-simulation-executor/tests/test_hf_access_validation.py index c1f409d44..0ef6f8b92 100644 --- a/projects/policyengine-simulation-executor/tests/test_hf_access_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_hf_access_validation.py @@ -5,12 +5,11 @@ from hashlib import sha256 import pytest - from policyengine_simulation_executor.hf_access_validation import ( EXPECTED_HF_TOKEN_DISPLAY_NAME, HF_RUNTIME_ENV_NAME, - HFArtifactAccess, HFAccessTarget, + HFArtifactAccess, build_uk_bundle_access_plan, validate_configured_uk_private_hf_access, validate_uk_private_hf_access, diff --git a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py index 2f4a71fe5..33837f67a 100644 --- a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py +++ b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py @@ -14,7 +14,6 @@ from types import SimpleNamespace import pytest - import src.modal._image_setup as image_setup from policyengine_simulation_executor.precompute_models import ArtifactManifest from src.modal._image_setup import fetch_artifacts diff --git a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py index c7444c3f6..e88c9755e 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py @@ -8,7 +8,6 @@ import subprocess import pytest - from fixtures.test_modal_scripts import REPO_ROOT, SCRIPTS_DIR pytest_plugins = ("fixtures.test_modal_scripts",) diff --git a/projects/policyengine-simulation-executor/tests/test_policyengine_dependency_source.py b/projects/policyengine-simulation-executor/tests/test_policyengine_dependency_source.py index 1fcd423c5..1450fe5b7 100644 --- a/projects/policyengine-simulation-executor/tests/test_policyengine_dependency_source.py +++ b/projects/policyengine-simulation-executor/tests/test_policyengine_dependency_source.py @@ -63,10 +63,7 @@ def test_all_runtime_dependencies_use_policyengine_models_extra(): for package in COMPONENT_PACKAGES: assert not any( isinstance(dependency, str) - and ( - dependency.startswith(f"{package}==") - or dependency.startswith(f"{package}[") - ) + and (dependency.startswith((f"{package}==", f"{package}["))) for dependency in dependencies ) diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index 07d5e4d40..eba242ee9 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -13,8 +13,6 @@ import pandas as pd import pytest -from pydantic import ValidationError - from fixtures.identity_stubs import install_identity_stubs from fixtures.wrapper_spm import wrapper_storage_id from policyengine_simulation_executor import precompute @@ -26,6 +24,7 @@ DatasetPlanEntry, PrecomputePlan, ) +from pydantic import ValidationError def _receipt() -> BundleVersionIdentity: @@ -354,7 +353,6 @@ def identity_stubs(self, monkeypatch): @pytest.fixture def fake_regions(self, monkeypatch): from policyengine.core.scoping_strategy import RowFilterStrategy - from policyengine_simulation_executor import simulation_runtime as sr def get_region(code): @@ -409,7 +407,6 @@ def test_dataset_filename_matches_runtime_stem_lookup(self, monkeypatch): dataset_logical_name, resolve_dataset_reference, ) - from policyengine_simulation_executor.artifact_keys import ( collect_dataset_identity, ) diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_precompute_app.py index 1ecfdd18b..0ff348119 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute_app.py @@ -12,9 +12,8 @@ import sys import pytest -from pydantic import ValidationError - from fixtures.fake_modal import install_fake_modal +from pydantic import ValidationError _WORKER_FUNCTIONS = ( "plan_artifacts", diff --git a/projects/policyengine-simulation-executor/tests/test_record_deployment.py b/projects/policyengine-simulation-executor/tests/test_record_deployment.py index 767e92dc9..fe976844d 100644 --- a/projects/policyengine-simulation-executor/tests/test_record_deployment.py +++ b/projects/policyengine-simulation-executor/tests/test_record_deployment.py @@ -8,7 +8,6 @@ from datetime import datetime import pytest - from policyengine_simulation_executor.precompute_models import DeployedMarker from src.modal.utils import record_deployment diff --git a/projects/policyengine-simulation-executor/tests/test_region_group_resolution.py b/projects/policyengine-simulation-executor/tests/test_region_group_resolution.py index 49982c466..d1f9f3122 100644 --- a/projects/policyengine-simulation-executor/tests/test_region_group_resolution.py +++ b/projects/policyengine-simulation-executor/tests/test_region_group_resolution.py @@ -8,7 +8,6 @@ RowFilterStrategy, WeightReplacementStrategy, ) - from policyengine_simulation_executor import simulation_runtime as sr from policyengine_simulation_executor.simulation_output_geographic import ( _should_build_us_congressional_district_impact as gate, diff --git a/projects/policyengine-simulation-executor/tests/test_release_bundle.py b/projects/policyengine-simulation-executor/tests/test_release_bundle.py index f03c53c9d..663e8dd6e 100644 --- a/projects/policyengine-simulation-executor/tests/test_release_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_release_bundle.py @@ -1,5 +1,6 @@ """Tests for policyengine.py release bundle helpers.""" +import json from copy import deepcopy from unittest.mock import MagicMock @@ -346,21 +347,20 @@ def test_resolve_runtime_bundle_dataset_uri_prefers_installed_default_dataset( dataset_path.write_bytes(b"data") receipt_path = tmp_path / BUNDLE_RECEIPT_FILENAME receipt_path.write_text( - """ - { - "bundle_version": "4.18.3", - "policyengine_version": "4.18.3", - "datasets": [ + json.dumps( { - "country": "us", - "dataset": "%s", - "version": "%s", - "path": "%s" + "bundle_version": "4.18.3", + "policyengine_version": "4.18.3", + "datasets": [ + { + "country": "us", + "dataset": bundle.default_dataset, + "version": bundle.data_version, + "path": str(dataset_path), + } + ], } - ] - } - """ - % (bundle.default_dataset, bundle.data_version, str(dataset_path)), + ), encoding="utf-8", ) monkeypatch.setenv("POLICYENGINE_BUNDLE_RECEIPT", str(receipt_path)) @@ -380,19 +380,18 @@ def test_resolve_runtime_bundle_dataset_uri_preserves_nondefault_override_with_r dataset_path.write_bytes(b"data") receipt_path = tmp_path / BUNDLE_RECEIPT_FILENAME receipt_path.write_text( - """ - { - "datasets": [ + json.dumps( { - "country": "us", - "dataset": "%s", - "version": "%s", - "path": "%s" + "datasets": [ + { + "country": "us", + "dataset": bundle.default_dataset, + "version": bundle.data_version, + "path": str(dataset_path), + } + ] } - ] - } - """ - % (bundle.default_dataset, bundle.data_version, str(dataset_path)), + ), encoding="utf-8", ) monkeypatch.setenv("POLICYENGINE_BUNDLE_RECEIPT", str(receipt_path)) diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 547c202e3..6155c8df8 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -5,7 +5,6 @@ import pytest from conftest import NoOpObservabilityRuntime from policyengine_simulation_contract.spm import SPMInputError - from policyengine_simulation_executor import simulation_runtime as sr from policyengine_simulation_executor import spm as executor_spm from policyengine_simulation_executor.national_partition import ( diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py b/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py index fa424eb31..59d2e3fa6 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py @@ -10,7 +10,6 @@ import numpy as np import pandas as pd import pytest - from policyengine_simulation_executor import segmented_national_reduce as reduce_mod from policyengine_simulation_executor.segmented_national_reduce import ( PrecomputedSimulation, diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py b/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py index bd249e401..5ff30aa04 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py @@ -1,6 +1,10 @@ """Contract tests for simulation API response shapes.""" -from policyengine_simulation_gateway.generate_openapi import create_openapi_app +from fixtures.test_simulation_api_contracts import ( + CURRENT_REQUIRED_BUDGET_KEYS, + CURRENT_SINGLE_YEAR_MACRO_KEYS, + CURRENT_SINGLE_YEAR_MACRO_RESULT, +) from policyengine_simulation_contract.gateway_models import ( BudgetWindowAnnualImpact, BudgetWindowResult, @@ -10,12 +14,7 @@ from policyengine_simulation_executor.simulation_macro_output import ( SingleYearMacroOutput, ) - -from fixtures.test_simulation_api_contracts import ( - CURRENT_REQUIRED_BUDGET_KEYS, - CURRENT_SINGLE_YEAR_MACRO_KEYS, - CURRENT_SINGLE_YEAR_MACRO_RESULT, -) +from policyengine_simulation_gateway.generate_openapi import create_openapi_app def test_job_status_result_preserves_current_single_year_macro_dict_contract(): diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_microdata.py b/projects/policyengine-simulation-executor/tests/test_simulation_microdata.py index 00e531814..64d913add 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_microdata.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_microdata.py @@ -4,7 +4,6 @@ import numpy as np import pandas as pd - from policyengine_simulation_executor.simulation_microdata import ( extract_output_microdata, rebuild_entity_frame, diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py index 69ff5c2fb..9774d7597 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py @@ -6,7 +6,6 @@ import pandas as pd from microdf import MicroDataFrame - from policyengine_simulation_executor.simulation_output_budget import ( build_budgetary_impact, ) diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 2e200af21..de27c4510 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -3,13 +3,12 @@ from __future__ import annotations import json -from contextlib import AbstractContextManager from collections import Counter +from contextlib import AbstractContextManager from types import SimpleNamespace import pandas as pd import pytest - from fixtures.spm_doubles import ( installed_spm_selection, spm_capable_simulation, @@ -29,41 +28,26 @@ FakeModelOutput, fake_analysis, ) -from policyengine_simulation_executor.release_bundle import BUNDLE_RECEIPT_FILENAME -from policyengine_simulation_executor.release_bundle import get_country_release_bundle +from policyengine_simulation_contract.uk_geography import ( + UKLocalAuthorityBoundaryVersion, + UKLocalAuthorityMetadata, +) from policyengine_simulation_executor.release_bundle import ( + BUNDLE_RECEIPT_FILENAME, + get_country_release_bundle, resolve_runtime_bundle_dataset_uri, ) -from policyengine_simulation_observability.stages import Stage -from policyengine_simulation_executor.simulation_runtime import ( - DatasetSelection, - RegionResolution, -) -from policyengine_simulation_executor.simulation_runtime import ( - _load_dataset as _load_selected_dataset, -) -from policyengine_simulation_executor.simulation_runtime import _nondefault_data_folder -from policyengine_simulation_executor.simulation_runtime import _normalise_policy -from policyengine_simulation_executor.simulation_runtime import ( - _resolve_dataset_reference, - _resolve_dataset_selection, -) -from policyengine_simulation_executor.simulation_runtime import _resolve_region -from policyengine_simulation_executor.simulation_runtime import ( - _run_simulation_impl_core, -) -from policyengine_simulation_executor.simulation_runtime import run_simulation_impl from policyengine_simulation_executor.simulation_macro_output import ( AgePovertyOutput, BaselineReformValue, BudgetaryImpact, BudgetaryOutput, + CongressionalDistrictImpactOutput, ConstituencyImpactRecord, DecileOutput, DetailedBudgetOutput, - CongressionalDistrictImpactOutput, - GeographicImpactOutput, GenderPovertyOutput, + GeographicImpactOutput, InequalityOutput, IntraDecileOutput, LocalAuthorityImpactRecord, @@ -72,17 +56,28 @@ PovertyOutput, SingleYearMacroOutput, ) +from policyengine_simulation_executor.simulation_output_builder import ( + SimulationOutputBuilder, +) from policyengine_simulation_executor.simulation_output_geographic import ( build_congressional_district_impact, build_congressional_district_impact_output, ) -from policyengine_simulation_executor.simulation_output_builder import ( - SimulationOutputBuilder, +from policyengine_simulation_executor.simulation_runtime import ( + DatasetSelection, + RegionResolution, + _nondefault_data_folder, + _normalise_policy, + _resolve_dataset_reference, + _resolve_dataset_selection, + _resolve_region, + _run_simulation_impl_core, + run_simulation_impl, ) -from policyengine_simulation_contract.uk_geography import ( - UKLocalAuthorityBoundaryVersion, - UKLocalAuthorityMetadata, +from policyengine_simulation_executor.simulation_runtime import ( + _load_dataset as _load_selected_dataset, ) +from policyengine_simulation_observability.stages import Stage def _params_with_spm(params, selection): diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py index d8a09cc9a..74b05af23 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py @@ -12,7 +12,6 @@ UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) - from policyengine_simulation_executor import simulation_output_geographic diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py b/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py index 68f9ed1d4..1c2ee7fb7 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py @@ -6,7 +6,6 @@ from policyengine.core import Dataset, Policy from policyengine.tax_benefit_models.common import reform as wrapper_reform from policyengine.utils.parametric_reforms import reform_dict_from_parameter_values - from policyengine_simulation_executor.simulation_runtime import ( DEFAULT_YEAR, DatasetSelection, diff --git a/projects/policyengine-simulation-executor/tests/test_smoke_app.py b/projects/policyengine-simulation-executor/tests/test_smoke_app.py index 1e37adc7a..ac7493958 100644 --- a/projects/policyengine-simulation-executor/tests/test_smoke_app.py +++ b/projects/policyengine-simulation-executor/tests/test_smoke_app.py @@ -4,13 +4,12 @@ import sys import pytest +from fixtures.fake_modal import install_fake_modal from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) -from fixtures.fake_modal import install_fake_modal - @pytest.fixture def smoke_module(monkeypatch): @@ -106,9 +105,11 @@ def test_uk_dataset_smoke_validates_installed_codes_against_packaged_resources( monkeypatch.setattr( smoke_module, "detect_uk_local_authority_metadata_from_hdf", - lambda path: observed.append(path) - or UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 + lambda path: ( + observed.append(path) + or UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 + ) ), ) diff --git a/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py b/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py index a43aee2f9..e7c78e8fa 100644 --- a/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py +++ b/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py @@ -51,8 +51,8 @@ def _worker_entrypoint(params, monkeypatch): def test_worker_year_boundary_matches_real_provider( forecast, period, unavailable_year, monkeypatch ): - from spm_calculator.policyengine_adapter import PolicyEngineSPMProvider from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail + from spm_calculator.policyengine_adapter import PolicyEngineSPMProvider provider = PolicyEngineSPMProvider(forecast, geography_kind="national") with pytest.raises(ValueError) as provider_error: @@ -78,8 +78,8 @@ def test_worker_year_boundary_matches_real_provider( def test_worker_scenario_boundary_matches_real_calculator_contract( forecast, period, monkeypatch ): - from spm_calculator.axiom_adapter import export_build_spec from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail + from spm_calculator.axiom_adapter import export_build_spec scenario = "scenario-not-in-the-pinned-forecast" with pytest.raises(ValueError) as adapter_error: @@ -100,8 +100,8 @@ def test_worker_scenario_boundary_matches_real_calculator_contract( def test_valid_worker_prevalidation_does_not_evaluate_amounts(forecast, monkeypatch): - from spm_calculator.policyengine_adapter import PolicyEngineSPMProvider from policyengine_simulation_executor.spm import normalize_runtime_spm + from spm_calculator.policyengine_adapter import PolicyEngineSPMProvider calculate = Mock(side_effect=AssertionError("Prevalidation is not measurement")) monkeypatch.setattr(PolicyEngineSPMProvider, "calculate_unit", calculate) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py index cacc5c332..b18164e29 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py @@ -12,18 +12,16 @@ import pandas as pd import pyarrow.parquet as pq import pytest -from policyengine_simulation_contract.stage12_execution import ( - UKLocalAuthorityBoundaryVersion, - UKLocalAuthorityMetadata, -) from policyengine_simulation_contract.spm import ( SPMRuntimeVersions, SPMSelection, build_spm_provenance, validate_spm_calculation_provenance, ) -from pydantic import ValidationError - +from policyengine_simulation_contract.stage12_execution import ( + UKLocalAuthorityBoundaryVersion, + UKLocalAuthorityMetadata, +) from policyengine_simulation_executor.stage12_artifacts import ( Stage12ArtifactStore, comparison_run_prefix, @@ -32,6 +30,7 @@ deserialize_uk_local_authority_metadata, serialize_simulation_frames, ) +from pydantic import ValidationError def _frames(): diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py index f92d3c3c4..21704607c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py @@ -8,7 +8,6 @@ import pytest from policyengine.bundle import get_current_bundle - from policyengine_simulation_executor.stage12_bundle import ( Stage12BundleError, assert_expected_bundle_values, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index ade681c12..811d90746 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -4,12 +4,11 @@ import importlib import sys -import tomllib from contextlib import contextmanager from pathlib import Path import pytest - +import tomllib from fixtures.fake_modal import install_fake_modal from policyengine_simulation_executor.release_bundle import ( get_bundled_package_version, 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 58ca4eef4..4c9b92301 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py @@ -22,14 +22,13 @@ SimulationRole, Stage12OutputPlan, ) -from pydantic import JsonValue -from policyengine_simulation_observability.stages import Stage - from policyengine_simulation_executor.stage12_runtime.output_planning import ( apply_output_plan, resolve_report_output_plan, validate_output_frames, ) +from policyengine_simulation_observability.stages import Stage +from pydantic import JsonValue EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_parity.py b/projects/policyengine-simulation-executor/tests/test_stage12_parity.py index ada90f10a..f5eb260e5 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_parity.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_parity.py @@ -3,7 +3,6 @@ import numpy as np import pandas as pd import pytest - from policyengine_simulation_executor.stage12_parity import ( NumericalTolerance, ParityMismatch, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py index ef215fae2..eae07ae75 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py @@ -7,7 +7,6 @@ import numpy as np import pandas as pd import pytest - from policyengine_simulation_contract.stage12_execution import SimulationRole from policyengine_simulation_executor.simulation_microdata import ( extract_output_microdata, @@ -17,7 +16,6 @@ qualify_report_parity, ) from policyengine_simulation_executor.stage12_runtime import SimulationCalculation - from test_stage12_runtime import _report diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py b/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py index 643bd3ba3..c75fc3fc6 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py @@ -4,7 +4,6 @@ from uuid import UUID import pytest - from policyengine_simulation_executor.stage12_result_comparison import compare_results RUN_ID = UUID("00000000-0000-0000-0000-000000000001") diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index f989c3a90..26b3251bb 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -23,6 +23,12 @@ StdoutLogDestination, configure, ) +from policyengine_simulation_contract.spm import ( + SPMRuntimeVersions, + SPMSelection, + build_spm_provenance, + validate_spm_calculation_provenance, +) from policyengine_simulation_contract.stage12_execution import ( ArtifactMediaType, ArtifactReference, @@ -51,12 +57,6 @@ Stage12OutputPlan, stage12_output_plan_sha256, ) -from policyengine_simulation_contract.spm import ( - SPMRuntimeVersions, - SPMSelection, - build_spm_provenance, - validate_spm_calculation_provenance, -) from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index a1abf01d9..057790992 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -6,7 +6,6 @@ from types import SimpleNamespace import pytest - from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle from policyengine_simulation_executor.stage12_worker_validation import ( _run_non_serving_calculation, diff --git a/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py b/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py index bba09b1bf..76869057e 100644 --- a/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py +++ b/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py @@ -4,11 +4,9 @@ from pathlib import Path from fastapi.testclient import TestClient - from fixtures.test_simulation_api_contracts import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_executor.main import app - PACKAGED_RUNTIME_MODULES = ( "policyengine_simulation_executor.compat_models", "policyengine_simulation_contract.hf_dataset", diff --git a/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py index 7bef14702..7fef5e41f 100644 --- a/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py +++ b/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py @@ -6,14 +6,12 @@ import pandas as pd import pytest - from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, ) - from policyengine_simulation_executor.uk_local_authority_metadata import ( - detect_uk_local_authority_metadata, detect_uk_local_authority_boundary_version, + detect_uk_local_authority_metadata, detect_uk_local_authority_metadata_from_hdf, load_uk_local_authority_resources, uk_area_regions_routed, diff --git a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py index c633f2930..dd1eb2f44 100644 --- a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py +++ b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py @@ -13,7 +13,6 @@ RowFilterStrategy, WeightReplacementStrategy, ) - from policyengine_simulation_executor import simulation_output_geographic from policyengine_simulation_executor import simulation_runtime as sr diff --git a/projects/policyengine-simulation-executor/tests/test_update_v2_version_manifest.py b/projects/policyengine-simulation-executor/tests/test_update_v2_version_manifest.py index 96a39c842..8a5df2d83 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_v2_version_manifest.py +++ b/projects/policyengine-simulation-executor/tests/test_update_v2_version_manifest.py @@ -5,7 +5,6 @@ from datetime import datetime, timezone import pytest - from policyengine_simulation_contract.stage12_manifest import V2WorkerValidation from src.modal.utils import update_v2_version_manifest as publisher diff --git a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index 9bd1544c0..56853459b 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -5,7 +5,6 @@ import sys import pytest - from src.modal.utils import update_version_registry as registry @@ -127,7 +126,6 @@ def test_bundle_manifest_records_policyengine_6_spm_capability( ): from policyengine.bundle import get_current_bundle from policyengine_simulation_contract.spm import SPM_CONTRACT_VERSION - from policyengine_simulation_executor import spm as executor_spm monkeypatch.setattr( diff --git a/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py b/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py index f3a521c9a..d6f37a66a 100644 --- a/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py @@ -3,7 +3,6 @@ from __future__ import annotations import pytest - from src.modal.utils import validate_v2_worker_release as validation diff --git a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py index 7caf5e078..f5ac439a0 100644 --- a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py +++ b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py @@ -243,7 +243,7 @@ def __init__(self, object_id: str = "mock-job-id-123"): def get(self, timeout: int = 0): if self.running: - raise TimeoutError() + raise TimeoutError if self.error is not None: raise self.error return self.result @@ -321,7 +321,6 @@ class OutputExpiredError(Exception): def mock_modal(monkeypatch): """Patch Modal calls in the gateway endpoints module.""" from policyengine_simulation_contract import budget_window_state, dataset_uri - from policyengine_simulation_gateway import endpoints mock_func = MockFunction() diff --git a/projects/policyengine-simulation-gateway/fixtures/package_imports.py b/projects/policyengine-simulation-gateway/fixtures/package_imports.py index e485fc66e..8c93aafa2 100644 --- a/projects/policyengine-simulation-gateway/fixtures/package_imports.py +++ b/projects/policyengine-simulation-gateway/fixtures/package_imports.py @@ -9,7 +9,6 @@ import pytest - 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/app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py index c15d4e867..2b4535033 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py @@ -8,10 +8,10 @@ and spawns jobs on those apps. """ -import modal import os from pathlib import Path +import modal from policyengine_simulation_observability.observability import ( modal_image_environment, ) @@ -75,18 +75,18 @@ def web_app(): GET /health - Health check """ from fastapi import FastAPI - from policyengine_simulation_observability.observability import ( init_simulation_observability, ) + from policyengine_simulation_gateway.auth import ( enforce_auth_configured_guard, enforce_production_auth_guard, ) - from policyengine_simulation_gateway.endpoints import router from policyengine_simulation_gateway.correlation import ( install_observability_id_middleware, ) + from policyengine_simulation_gateway.endpoints import router api = FastAPI( title="PolicyEngine Simulation Gateway", diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/auth.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/auth.py index c0ddec199..bd4265f28 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/auth.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/auth.py @@ -33,9 +33,8 @@ from fastapi import Depends, HTTPException, Request, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer - -from policyengine_observability import ObservabilityRuntime from policyengine_fastapi.auth import JWTDecoder +from policyengine_observability import ObservabilityRuntime logger = logging.getLogger(__name__) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py index ae3181740..7fa67f982 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py @@ -13,7 +13,6 @@ from pathlib import Path from fastapi import FastAPI - from policyengine_simulation_contract.gateway_models import ( BudgetWindowBatchRequest, BudgetWindowBatchStatusResponse, 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 9b653fc7e..e9c2c5f3d 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py @@ -3,7 +3,6 @@ from __future__ import annotations from fastapi.responses import JSONResponse - from policyengine_simulation_contract.gateway_models import ( BudgetWindowBatchStatusResponse, ) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py index b12c16102..aecb829c2 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py @@ -27,9 +27,10 @@ def smoke_import_gateway() -> dict: import pkgutil import policyengine_simulation_contract - import policyengine_simulation_gateway import policyengine_simulation_observability + import policyengine_simulation_gateway + imported = [] for package in ( policyengine_simulation_gateway, @@ -53,7 +54,7 @@ def smoke_import_gateway() -> dict: return { "modules_imported": len(imported), "routes": sorted( - f"{sorted(route.methods)[0]} {route.path}" + f"{min(route.methods)} {route.path}" for route in asgi_app.routes if hasattr(route, "methods") ), diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/testing.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/testing.py index 8648a89b1..2b1e27d13 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/testing.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/testing.py @@ -6,10 +6,10 @@ """ from fastapi import FastAPI - from policyengine_simulation_observability.observability import ( init_simulation_observability, ) + from policyengine_simulation_gateway.auth import require_auth from policyengine_simulation_gateway.correlation import ( install_observability_id_middleware, diff --git a/projects/policyengine-simulation-gateway/tests/conftest.py b/projects/policyengine-simulation-gateway/tests/conftest.py index 4ad5e8526..71f36ff92 100644 --- a/projects/policyengine-simulation-gateway/tests/conftest.py +++ b/projects/policyengine-simulation-gateway/tests/conftest.py @@ -1,12 +1,10 @@ """Root pytest configuration for gateway tests.""" -import importlib import sys from pathlib import Path import pytest from fastapi.testclient import TestClient - from policyengine_simulation_gateway.testing import create_gateway_app project_root = Path(__file__).parent.parent diff --git a/projects/policyengine-simulation-gateway/tests/test_auth.py b/projects/policyengine-simulation-gateway/tests/test_auth.py index 3639823b2..906bbcae5 100644 --- a/projects/policyengine-simulation-gateway/tests/test_auth.py +++ b/projects/policyengine-simulation-gateway/tests/test_auth.py @@ -2,14 +2,13 @@ from __future__ import annotations -import pytest -from fastapi.testclient import TestClient from types import SimpleNamespace from unittest.mock import Mock -from policyengine_simulation_gateway.testing import create_gateway_app +import pytest +from fastapi.testclient import TestClient from policyengine_simulation_gateway import auth as auth_module - +from policyengine_simulation_gateway.testing import create_gateway_app GATED_REQUESTS = [ ( @@ -185,8 +184,8 @@ def test__given_repeated_requests__then_decoder_not_reinstantiated(monkeypatch): decoder. We spy on ``JWTDecoder.__init__`` and assert it runs at most once across many requests.""" - from policyengine_simulation_gateway.testing import create_gateway_app from policyengine_fastapi.auth import jwt_decoder as jwt_decoder_module + from policyengine_simulation_gateway.testing import create_gateway_app monkeypatch.delenv(auth_module.GATEWAY_AUTH_DISABLED_ENV, raising=False) monkeypatch.setenv(auth_module.GATEWAY_AUTH_REQUIRED_ENV, "1") diff --git a/projects/policyengine-simulation-gateway/tests/test_endpoints.py b/projects/policyengine-simulation-gateway/tests/test_endpoints.py index 098174be5..5dad321f4 100644 --- a/projects/policyengine-simulation-gateway/tests/test_endpoints.py +++ b/projects/policyengine-simulation-gateway/tests/test_endpoints.py @@ -10,14 +10,13 @@ import pytest from fastapi.testclient import TestClient -from policyengine_simulation_contract.hf_dataset import HuggingFaceDatasetReferenceError -from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER - from fixtures.gateway_endpoints import ( TEST_APP_RELEASE_BUNDLE, TEST_ROUTING_STATE, resolve_test_dataset_uri, ) +from policyengine_simulation_contract.hf_dataset import HuggingFaceDatasetReferenceError +from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000001" diff --git a/projects/policyengine-simulation-gateway/tests/test_package_imports.py b/projects/policyengine-simulation-gateway/tests/test_package_imports.py index ceb1457a3..6c1cfd575 100644 --- a/projects/policyengine-simulation-gateway/tests/test_package_imports.py +++ b/projects/policyengine-simulation-gateway/tests/test_package_imports.py @@ -1,5 +1,5 @@ -import sys import importlib +import sys def test_gateway_models_import_does_not_import_fastapi_endpoints( diff --git a/projects/policyengine-simulation-gateway/tests/test_route_table.py b/projects/policyengine-simulation-gateway/tests/test_route_table.py index 233c673c0..1e9fcd098 100644 --- a/projects/policyengine-simulation-gateway/tests/test_route_table.py +++ b/projects/policyengine-simulation-gateway/tests/test_route_table.py @@ -5,7 +5,6 @@ """ from fastapi.routing import APIRoute - from policyengine_simulation_gateway.generate_openapi import create_openapi_app from policyengine_simulation_gateway.testing import create_gateway_app diff --git a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py index 53d214880..18c73de46 100644 --- a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py +++ b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py @@ -3,10 +3,8 @@ from copy import deepcopy import pytest - from fixtures.gateway_endpoints import TEST_ROUTING_STATE - CAPABILITY = { "contract_version": "canonical-spm-v1", "defaults": {"forecast_content_sha256": "a" * 64, "scenario": "ce_trend"}, diff --git a/ruff.toml b/ruff.toml index 21a4f7848..482af947b 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,4 +1,14 @@ -# Ruff's defaults as of 0.15, made explicit so a ruff upgrade cannot change -# which rules apply. Widen this deliberately, one rule family per PR. +# Rule families are enabled one themed group per PR, each fixing everything it +# finds (no ratchet), until the select list is "ALL". See +# docs/engineering/skills/rules-registry.md and rules/registry.toml. +# There is no root pyproject.toml, so files outside the projects and libs would +# be treated as Python 3.9. scripts/rules supports Python 3.11+ (and is shared +# byte-for-byte with policyengine-api, which targets 3.11). +per-file-target-version = { "scripts/**" = "py311" } + [lint] -select = ["E4", "E7", "E9", "F"] +select = [ + "E4", "E7", "E9", "F", + # Group 1: safe auto-fixes. + "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W", +] diff --git a/rules/registry.toml b/rules/registry.toml index 9df1b8712..e6883113c 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -611,16 +611,20 @@ kind = "none" [[rule]] id = "sim.quality.ruff-lint" category = "quality" -statement = "Python code passes ruff check with the rules selected in ruff.toml (ruff's defaults, made explicit); CI does not run ruff check yet." -rationale = "docs/engineering/skills/testing.md#common-commands" -scope = ["projects/*/src/**/*.py", "libs/*/src/**/*.py"] +statement = "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." +rationale = "docs/engineering/skills/rules-registry.md#baselines" +scope = ["projects/*/src/**/*.py", "libs/*/src/**/*.py", "projects/*/tests/**/*.py", "libs/*/tests/**/*.py"] tier = "mechanical" -status = "advisory" +status = "enforced" autofix = true [rule.enforcer] kind = "ruff" mode = "lint" -codes = ["E4", "E7", "E9", "F"] +codes = ["E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W"] +[[rule.ci]] +workflow = ".github/workflows/pr.yml" +job = "lint" +invocation = "ruff check ." [[rule]] id = "sim.quality.tests-pass" diff --git a/scripts/rules/tests/support.py b/scripts/rules/tests/support.py index d10653122..3d89aebfb 100644 --- a/scripts/rules/tests/support.py +++ b/scripts/rules/tests/support.py @@ -6,6 +6,7 @@ import textwrap from datetime import date from pathlib import Path +from typing import Self from scripts.rules import TOOL_VERSION from scripts.rules.guard import GuardResult, run_guard @@ -197,7 +198,7 @@ def __init__(self, files: dict[str, str] | None = None, *, render: bool = True): if render: self.render() - def __enter__(self) -> Repo: + def __enter__(self) -> Self: return self def __exit__(self, *exc: object) -> None: diff --git a/scripts/run_quality_guards.py b/scripts/run_quality_guards.py index d347fb40e..c3d8526ff 100644 --- a/scripts/run_quality_guards.py +++ b/scripts/run_quality_guards.py @@ -6,7 +6,6 @@ from collections.abc import Callable from pathlib import Path - REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) @@ -14,7 +13,6 @@ from scripts.rules import ci_scripts # noqa: E402 from scripts.rules import guard as rules_guard # noqa: E402 - Guard = tuple[str, Callable[[], list[str]]] GUARDS: tuple[Guard, ...] = (