From 16f112a6d74c08d3d538240830643a3baeb3629f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:42:23 +0400 Subject: [PATCH 01/12] Apply ruff's safe docstring fixes Run ruff check --fix for the pydocstyle (D) rules: put summaries on the first line, drop blank lines after function docstrings, and normalise blank lines around sections and closing quotes. Only docstring whitespace changes; the published OpenAPI documents and Pydantic JSON schemas are unchanged. Co-Authored-By: Claude Opus 5.5 --- .github/find-api-model-versions.py | 8 +- .github/scripts/check_updates.py | 3 +- .../scripts/verify_observability_collector.py | 1 - migrations/v2/env.py | 1 - policyengine_api/api.py | 4 +- policyengine_api/asgi_factory.py | 2 - policyengine_api/constants.py | 3 +- policyengine_api/country.py | 5 +- .../data/congressional_districts.py | 38 +++++----- policyengine_api/data/orm.py | 5 -- policyengine_api/data/places.py | 11 ++- .../data/v2/catalog/catalog_selection.py | 2 - .../data/v2/catalog/extraction.py | 3 - .../data/v2/catalog/initialization.py | 3 - .../data/v2/catalog/publication.py | 1 - .../data/v2/catalog/publication_staging.py | 1 - policyengine_api/data/v2/catalog/records.py | 6 -- policyengine_api/data/v2/database.py | 7 -- .../v2/household_migration_qualification.py | 6 -- .../data/v2/metadata_validation.py | 1 - policyengine_api/data/v2/migration_target.py | 3 - policyengine_api/data/v2/models/base.py | 2 - .../data/v2/policy_migration_qualification.py | 6 -- policyengine_api/data/v2/report_runs.py | 6 -- policyengine_api/data/v2/settings.py | 9 --- policyengine_api/dataset_display.py | 1 - policyengine_api/fastapi_routes/responses.py | 1 - policyengine_api/fastapi_routes/v2/errors.py | 1 - .../fastapi_routes/v2/households/routes.py | 2 - .../fastapi_routes/v2/policies/routes.py | 2 - policyengine_api/fastapi_routes/v2/routes.py | 1 - .../v2/user_households/routes.py | 1 - .../fastapi_routes/v2/user_policies/routes.py | 1 - policyengine_api/gcp_logging.py | 1 - .../libs/simulation_entrypoint.py | 44 +++++------ policyengine_api/migration_flags.py | 4 - policyengine_api/migration_logging.py | 1 - policyengine_api/observability/identifiers.py | 2 - policyengine_api/observability/runtime.py | 1 - policyengine_api/query_parameters.py | 6 -- policyengine_api/readiness.py | 1 - policyengine_api/request_context.py | 7 -- policyengine_api/routes/household_routes.py | 9 +-- policyengine_api/routes/metadata_routes.py | 2 +- policyengine_api/routes/policy_routes.py | 12 +-- .../routes/reform_impact_routes.py | 1 - .../routes/report_output_routes.py | 14 ++-- policyengine_api/routes/simulation_routes.py | 14 ++-- .../routes/user_profile_routes.py | 9 +-- policyengine_api/runtime_cache/client.py | 1 - policyengine_api/runtime_cache/core.py | 2 - .../runtime_cache/reform_impacts.py | 4 - policyengine_api/runtime_cache/settings.py | 1 - .../services/budget_window_cache.py | 1 - .../services/economy_comparison.py | 32 ++++---- policyengine_api/services/economy_service.py | 75 +++++-------------- .../services/household_calculation_service.py | 4 - .../services/household_mirroring.py | 1 - .../services/household_service.py | 1 - policyengine_api/services/policy_mirroring.py | 1 - policyengine_api/services/policy_service.py | 2 - .../services/reform_impacts_service.py | 3 - .../services/report_output_service.py | 2 - .../services/user_policy_mirroring.py | 2 - .../services/user_policy_service.py | 1 - .../services/v2/comparison_runs/validators.py | 3 - .../services/v2/households/validators.py | 3 - .../services/v2/metadata/validators.py | 1 - .../services/v2/policies/validators.py | 4 - policyengine_api/services/v2/reports/types.py | 1 - .../v2/user_policies/transformations.py | 2 - policyengine_api/spm.py | 5 -- policyengine_api/utils/cache_utils.py | 2 +- policyengine_api/utils/input_validation.py | 1 - .../payload_validators/validate_country.py | 1 + .../validate_household_payload.py | 4 +- scripts/export_migration_contracts.py | 4 - scripts/process_v1_household_mirror_event.py | 3 - scripts/rules/ci_scripts.py | 3 +- scripts/v1_database_migration.py | 2 - tests/conftest.py | 1 - .../test_environment_variables.py | 2 - tests/fixtures/integration/simulations.py | 4 +- tests/fixtures/libs/simulation_entrypoint.py | 10 +-- tests/fixtures/services/economy_service.py | 4 +- tests/fixtures/v2_catalog.py | 4 - tests/integration/conftest.py | 2 - .../test_alembic_mysql_lifecycle.py | 1 - tests/integration/test_live_v2_policies.py | 1 - tests/integration/test_runtime_cache_redis.py | 3 +- .../test_v1_schema_metadata_compatibility.py | 1 - tests/to_refactor/conftest.py | 1 - .../python/test_validate_country.py | 6 +- tests/unit/conftest.py | 1 - tests/unit/libs/test_simulation_entrypoint.py | 3 +- .../routes/test_metadata_routes_streaming.py | 1 - .../unit/services/test_economy_comparison.py | 21 ++++-- tests/unit/services/test_metadata_service.py | 7 +- tests/unit/test_country_spm.py | 3 +- tests/unit/utils/test_cache_utils.py | 3 +- 100 files changed, 144 insertions(+), 382 deletions(-) diff --git a/.github/find-api-model-versions.py b/.github/find-api-model-versions.py index 0c0096db2..55287ca08 100644 --- a/.github/find-api-model-versions.py +++ b/.github/find-api-model-versions.py @@ -48,9 +48,7 @@ def _data_release_version(bundle: dict, country_id: str) -> str: def find_api_model_versions() -> dict[str, str]: - """ - Find and validate the package and data versions in the installed bundle. - """ + """Find and validate the package and data versions in the installed bundle.""" bundle = get_current_bundle() if not isinstance(bundle, dict): raise RuntimeError("Installed PolicyEngine bundle manifest is not an object.") @@ -89,9 +87,7 @@ def find_api_model_versions() -> dict[str, str]: def find_api_model_versions_and_output_to_github() -> None: - """ - Find the API model versions and output them to a file for GitHub. - """ + """Find the API model versions and output them to a file for GitHub.""" versions = find_api_model_versions() with open(os.environ["GITHUB_ENV"], "a") as f: for key, value in versions.items(): diff --git a/.github/scripts/check_updates.py b/.github/scripts/check_updates.py index eccbc982e..4d59de47b 100644 --- a/.github/scripts/check_updates.py +++ b/.github/scripts/check_updates.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 -""" -Fetch and format changelog entries between two versions of a package. +"""Fetch and format changelog entries between two versions of a package. Fetches CHANGELOG.md from the package's GitHub repo and extracts entries between the old and new version, formatted as markdown. diff --git a/.github/scripts/verify_observability_collector.py b/.github/scripts/verify_observability_collector.py index 8156a70bc..7d387a93c 100644 --- a/.github/scripts/verify_observability_collector.py +++ b/.github/scripts/verify_observability_collector.py @@ -315,7 +315,6 @@ def verify( delivery_timeout_seconds: float = 600.0, ) -> VerificationProbe: """Export real telemetry, verify storage, and require log rejection.""" - started_at = time.time() verification_id = f"collector-verify-{uuid.uuid4().hex}" trace_request, metric_request, log_request, trace_id = _probe_requests( diff --git a/migrations/v2/env.py b/migrations/v2/env.py index 17b8e7d6d..cbde16f98 100644 --- a/migrations/v2/env.py +++ b/migrations/v2/env.py @@ -38,7 +38,6 @@ def _include_application_object( _compare_to: SchemaItem | None, ) -> bool: """Keep Alembic's own version table outside application drift.""" - return not (type_ == "table" and name == "alembic_version") diff --git a/policyengine_api/api.py b/policyengine_api/api.py index 81fc92106..ed2834da8 100644 --- a/policyengine_api/api.py +++ b/policyengine_api/api.py @@ -1,7 +1,5 @@ # ruff: noqa: E402 -""" -This is the main Flask app for the PolicyEngine API. -""" +"""This is the main Flask app for the PolicyEngine API.""" import sys import time diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py index 0c15fb9d6..7b0c91354 100644 --- a/policyengine_api/asgi_factory.py +++ b/policyengine_api/asgi_factory.py @@ -96,7 +96,6 @@ def _matched_route_template(routes: Iterable[BaseRoute], scope: Scope) -> str | def _native_route_template(app: FastAPI, scope: Scope) -> str | None: """Return the matched FastAPI route template, excluding the Flask mount.""" - return _matched_route_template(app.router.routes, scope) @@ -109,7 +108,6 @@ def create_asgi_app( observability_runtime: ObservabilityRuntime | None = None, ) -> ASGIApp: """Create the Stage 2 FastAPI shell around the existing Flask app.""" - if route_settings is None: route_settings = RouteImplementationSettings.from_environment() if dependencies is None: diff --git a/policyengine_api/constants.py b/policyengine_api/constants.py index cea6ecd7d..bcff670e5 100644 --- a/policyengine_api/constants.py +++ b/policyengine_api/constants.py @@ -185,8 +185,7 @@ def get_bundle_default_dataset_option(country_id: str) -> dict: def _build_runtime_cache_version( scope: str, country_id: str, caller_version: str | None = None ) -> str: - """ - Build a compact version token for cache keys stored in legacy VARCHAR(10) + """Build a compact version token for cache keys stored in legacy VARCHAR(10) columns. The token changes whenever the relevant runtime or payload schema changes, even if the country package version is unchanged. """ diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 332871fae..66c6ff415 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -557,7 +557,6 @@ def prepare_calculation( spm_requested: bool = False, ) -> PreparedCountryCalculation: """Parse a situation before the API accepts it as a calculation.""" - simulation, system = self._create_simulation(household, reform, spm=spm) household = json.loads(json.dumps(household)) has_axes = "axes" in household @@ -648,14 +647,14 @@ def _apply_reform_to_system( def create_policy_reform(policy_data: dict) -> dict: - """ - Create a policy reform. + """Create a policy reform. Args: policy_data (dict): The policy data. Returns: dict: The reform. + """ def modify_parameters(parameters: ParameterNode) -> ParameterNode: diff --git a/policyengine_api/data/congressional_districts.py b/policyengine_api/data/congressional_districts.py index c78752304..38a826cc7 100644 --- a/policyengine_api/data/congressional_districts.py +++ b/policyengine_api/data/congressional_districts.py @@ -1,5 +1,4 @@ -""" -US Congressional District Metadata +"""US Congressional District Metadata This module defines the metadata for all 435 US Congressional districts based on the 2020 Census apportionment, effective for the 118th Congress @@ -68,8 +67,7 @@ class CongressionalDistrictMetadataItem(BaseModel): - """ - Metadata for a single US Congressional district. + """Metadata for a single US Congressional district. Uses Pydantic BaseModel for: - Runtime validation of data integrity @@ -588,8 +586,7 @@ class CongressionalDistrictMetadataItem(BaseModel): def _get_ordinal_suffix(number: int) -> str: - """ - Get the ordinal suffix for a number (st, nd, rd, th). + """Get the ordinal suffix for a number (st, nd, rd, th). Examples: 1 -> "st" @@ -599,6 +596,7 @@ def _get_ordinal_suffix(number: int) -> str: 11 -> "th" 21 -> "st" 22 -> "nd" + """ if 10 <= number % 100 <= 20: # Special case for 11th, 12th, 13th, etc. @@ -609,33 +607,32 @@ def _get_ordinal_suffix(number: int) -> str: def _format_district_number(number: int) -> str: - """ - Format district number with leading zero for single digits. + """Format district number with leading zero for single digits. Examples: 1 -> "01" 9 -> "09" 10 -> "10" 38 -> "38" + """ return f"{number:02d}" def _build_district_name(state_code: str, number: int) -> str: - """ - Build the district name in the format: congressional_district/- + """Build the district name in the format: congressional_district/- Examples: ("CA", 5) -> "congressional_district/CA-05" ("TX", 38) -> "congressional_district/TX-38" ("DC", 1) -> "congressional_district/DC-01" + """ return f"congressional_district/{state_code}-{_format_district_number(number)}" def _build_district_label(state_code: str, number: int) -> str: - """ - Build the district label in the format: 's th congressional district + """Build the district label in the format: 's th congressional district For at-large districts (states with only 1 district), use: 's at-large congressional district Examples: @@ -645,6 +642,7 @@ def _build_district_label(state_code: str, number: int) -> str: ("FL", 21) -> "Florida's 21st congressional district" ("AK", 1) -> "Alaska's at-large congressional district" ("WY", 1) -> "Wyoming's at-large congressional district" + """ state_name = STATE_CODE_TO_NAME[state_code] if state_code in AT_LARGE_STATES: @@ -654,8 +652,7 @@ def _build_district_label(state_code: str, number: int) -> str: def build_congressional_district_metadata() -> list[dict]: - """ - Build the complete congressional district metadata structure for use in country.py. + """Build the complete congressional district metadata structure for use in country.py. Returns a list of dictionaries with 'name', 'label', 'type', 'state_abbreviation', and 'state_name' keys, formatted as: @@ -679,6 +676,7 @@ def build_congressional_district_metadata() -> list[dict]: Returns: List of 436 dictionaries (435 voting districts + DC) + """ return [ { @@ -693,23 +691,23 @@ def build_congressional_district_metadata() -> list[dict]: def get_valid_state_codes() -> set[str]: - """ - Get the set of valid US state codes (lowercase for case-insensitive matching). + """Get the set of valid US state codes (lowercase for case-insensitive matching). Returns: Set of 51 lowercase state codes (50 states + DC) + """ return {code.lower() for code in STATE_CODE_TO_NAME.keys()} def get_valid_congressional_districts() -> set[str]: - """ - Get the set of valid congressional district identifiers (lowercase for case-insensitive matching). + """Get the set of valid congressional district identifiers (lowercase for case-insensitive matching). Format: "-" (e.g., "ca-37", "tx-01") Returns: Set of 436 lowercase district identifiers + """ return { f"{district.state_code.lower()}-{_format_district_number(district.number)}" @@ -718,8 +716,7 @@ def get_valid_congressional_districts() -> set[str]: def normalize_us_region(region: str) -> str: - """ - Normalize a US region string to the standard prefixed format. + """Normalize a US region string to the standard prefixed format. This function handles legacy region formats (bare state codes like "ca") and converts them to the standard format ("state/ca"). It should be called @@ -739,6 +736,7 @@ def normalize_us_region(region: str) -> str: Note: This function does NOT validate that the region is valid - it only normalizes the format. Use _validate_us_region for validation. + """ # Already has a valid prefix - return as-is if region.startswith(("state/", "place/", "congressional_district/")): diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 1a1ffaa67..57f345e2d 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -40,7 +40,6 @@ def get_remote_database_config() -> dict[str, str]: def build_session_factory(engine: Engine) -> sessionmaker[Session]: """Return SQLAlchemy's standard configurable Session factory.""" - return sessionmaker(bind=engine, class_=Session, expire_on_commit=False) @@ -77,7 +76,6 @@ def get_connection() -> DBAPIConnection: def get_v1_engine() -> Engine: """Return the process-owned Cloud SQL engine; never select SQLite.""" - global _v1_engine if _v1_engine is None: _v1_engine = _build_remote_engine() @@ -86,7 +84,6 @@ def get_v1_engine() -> Engine: def get_v1_session_factory() -> sessionmaker[Session]: """Return one configured Session factory for the Cloud SQL engine.""" - global _v1_session_factory if _v1_session_factory is None: _v1_session_factory = build_session_factory(get_v1_engine()) @@ -95,14 +92,12 @@ def get_v1_session_factory() -> sessionmaker[Session]: def clear_v1_session_factories() -> None: """Forget cached factories, primarily after replacing an Engine in tests.""" - global _v1_session_factory _v1_session_factory = None def close_v1_engines() -> None: """Release process-owned SQLAlchemy pools and Cloud SQL connectors.""" - global _cloud_sql_connector, _v1_engine clear_v1_session_factories() if _v1_engine is not None: diff --git a/policyengine_api/data/places.py b/policyengine_api/data/places.py index f24467f83..3a73b950b 100644 --- a/policyengine_api/data/places.py +++ b/policyengine_api/data/places.py @@ -1,5 +1,4 @@ -""" -US Census place code parsing and validation utilities. +"""US Census place code parsing and validation utilities. Place codes follow the format: STATE_ABBREV-PLACE_FIPS Example: NJ-57000 for Newark, NJ @@ -9,8 +8,7 @@ def parse_place_code(place_code: str) -> tuple[str, str]: - """ - Parse a place code into its state abbreviation and FIPS components. + """Parse a place code into its state abbreviation and FIPS components. Args: place_code: Place code in format STATE_ABBREV-PLACE_FIPS (e.g., "NJ-57000") @@ -20,6 +18,7 @@ def parse_place_code(place_code: str) -> tuple[str, str]: Raises: ValueError: If the place code format is invalid + """ if "-" not in place_code: raise ValueError( @@ -30,14 +29,14 @@ def parse_place_code(place_code: str) -> tuple[str, str]: def validate_place_code(place_code: str) -> None: - """ - Validate a place code has valid state abbreviation and FIPS format. + """Validate a place code has valid state abbreviation and FIPS format. Args: place_code: Place code in format STATE_ABBREV-PLACE_FIPS (e.g., "NJ-57000") Raises: ValueError: If the state abbreviation or FIPS code is invalid + """ state_abbrev, place_fips = parse_place_code(place_code) diff --git a/policyengine_api/data/v2/catalog/catalog_selection.py b/policyengine_api/data/v2/catalog/catalog_selection.py index f998cf224..f72a43192 100644 --- a/policyengine_api/data/v2/catalog/catalog_selection.py +++ b/policyengine_api/data/v2/catalog/catalog_selection.py @@ -44,7 +44,6 @@ class SelectedCatalog: def validate_policyengine_version(value: str) -> str: """Return one bounded canonical PEP 440 version string.""" - if not isinstance(value, str) or not value or value != value.strip(): raise InvalidPolicyEngineVersionError( "policyengine_version must be a non-empty canonical version" @@ -74,7 +73,6 @@ def select_catalog( policyengine_version: str | None = None, ) -> SelectedCatalog: """Select one country catalog using the requested or running package version.""" - if country_id not in SUPPORTED_V2_COUNTRY_IDS: raise UnsupportedPreviewCountryError(country_id) explicit_version = policyengine_version is not None diff --git a/policyengine_api/data/v2/catalog/extraction.py b/policyengine_api/data/v2/catalog/extraction.py index 421d4633b..4d8a24fc8 100644 --- a/policyengine_api/data/v2/catalog/extraction.py +++ b/policyengine_api/data/v2/catalog/extraction.py @@ -142,7 +142,6 @@ def _type_name(value: object) -> str | None: def normalize_json_value(value: object) -> JSONValue: """Convert supported public values into strict JSON-compatible values.""" - if value is None or isinstance(value, (str, bool, int)): return value if isinstance(value, float): @@ -718,7 +717,6 @@ def extract_catalog( installed_version: Callable[[str], str] = importlib_metadata.version, ) -> NormalizedCatalog: """Validate and normalize supplied PolicyEngine.py public objects.""" - dependencies = _verify_bundle( bundle=bundle, policyengine_version=policyengine_version, @@ -751,7 +749,6 @@ def extract_catalog( def extract_installed_catalog() -> NormalizedCatalog: """Load and normalize the installed PolicyEngine.py certified catalog.""" - import policyengine from policyengine.bundle import get_current_bundle diff --git a/policyengine_api/data/v2/catalog/initialization.py b/policyengine_api/data/v2/catalog/initialization.py index 556519694..9fcf31c54 100644 --- a/policyengine_api/data/v2/catalog/initialization.py +++ b/policyengine_api/data/v2/catalog/initialization.py @@ -32,7 +32,6 @@ def build_data_write_engine(settings: V2DatabaseSettings) -> Engine: """Build an isolated, unpooled engine for one initialization attempt.""" - return create_engine(settings.connection.url, poolclass=NullPool) @@ -46,7 +45,6 @@ def initialize_catalog( ), ) -> PublicationEvidence: """Load only row-write settings, extract fully, then publish explicitly.""" - settings = load_v2_data_write_database_settings(environ) catalog = extractor() engine = engine_builder(settings) @@ -78,7 +76,6 @@ def _error_payload(error: Exception) -> dict[str, object]: def main() -> int: """Run explicit initialization and return a shell-compatible status.""" - try: evidence = initialize_catalog() except Exception as error: # noqa: BLE001 - command must return safe evidence diff --git a/policyengine_api/data/v2/catalog/publication.py b/policyengine_api/data/v2/catalog/publication.py index 8f33fc664..cb97ed763 100644 --- a/policyengine_api/data/v2/catalog/publication.py +++ b/policyengine_api/data/v2/catalog/publication.py @@ -93,7 +93,6 @@ def publish_catalog( checkpoint: Callable[[str, Connection], None] | None = None, ) -> PublicationEvidence: """Publish one complete catalog atomically and return non-secret evidence.""" - started = time.monotonic() with engine.begin() as connection: _verify_expected_revision(connection) diff --git a/policyengine_api/data/v2/catalog/publication_staging.py b/policyengine_api/data/v2/catalog/publication_staging.py index d45dcb5f6..2d2f556b0 100644 --- a/policyengine_api/data/v2/catalog/publication_staging.py +++ b/policyengine_api/data/v2/catalog/publication_staging.py @@ -275,7 +275,6 @@ def copy_rows( rows: Iterable[Sequence[object]], ) -> int: """Write one bounded source stream through Psycopg COPY.""" - raw_connection = connection.connection.driver_connection statement = sql.SQL("COPY {} ({}) FROM STDIN").format( sql.Identifier(table.name), diff --git a/policyengine_api/data/v2/catalog/records.py b/policyengine_api/data/v2/catalog/records.py index 6a0adf232..b8a5864fe 100644 --- a/policyengine_api/data/v2/catalog/records.py +++ b/policyengine_api/data/v2/catalog/records.py @@ -117,7 +117,6 @@ def iter_batches( batch_size: int, ) -> Iterator[tuple[RecordT, ...]]: """Yield bounded immutable slices without constructing a flattened copy.""" - if batch_size < 1: raise ValueError("batch_size must be positive") for start in range(0, len(records), batch_size): @@ -142,7 +141,6 @@ def parameter_value_batches( batch_size: int, ) -> Iterator[tuple[ParameterValueRecord, ...]]: """Yield canonical parameter values in bounded batches.""" - if batch_size < 1: raise ValueError("batch_size must be positive") batch: list[ParameterValueRecord] = [] @@ -157,7 +155,6 @@ def parameter_value_batches( def entity_counts(self) -> dict[str, int]: """Return non-secret record counts for validation and deployment evidence.""" - return { "models": 1, "model_versions": 1, @@ -180,7 +177,6 @@ class NormalizedCatalog: def country(self, country_id: str) -> CountryCatalog: """Return one supported country catalog.""" - for catalog in self.countries: if catalog.country_id == country_id: return catalog @@ -188,7 +184,6 @@ def country(self, country_id: str) -> CountryCatalog: def entity_counts(self) -> dict[str, int]: """Return aggregate counts without catalog content.""" - counts: dict[str, int] = {} for country in self.countries: for name, count in country.entity_counts().items(): @@ -198,5 +193,4 @@ def entity_counts(self) -> dict[str, int]: def record_content(record: object) -> tuple[Any, ...]: """Return a deterministic field-ordered representation for comparisons.""" - return tuple(getattr(record, item.name) for item in fields(record)) diff --git a/policyengine_api/data/v2/database.py b/policyengine_api/data/v2/database.py index b742177bd..0f93f9b16 100644 --- a/policyengine_api/data/v2/database.py +++ b/policyengine_api/data/v2/database.py @@ -33,7 +33,6 @@ def _url_fingerprint(settings: V2DatabaseSettings) -> tuple[object, ...]: """Build a private comparison value that is never rendered or logged.""" - url = settings.connection.url return ( url.drivername, @@ -50,7 +49,6 @@ def _url_fingerprint(settings: V2DatabaseSettings) -> tuple[object, ...]: def build_v2_engine(settings: V2DatabaseSettings) -> Engine: """Construct an unconnected SQLModel engine for an explicit v2 target.""" - return create_engine( settings.connection.url, pool_pre_ping=True, @@ -67,13 +65,11 @@ def build_v2_engine(settings: V2DatabaseSettings) -> Engine: def build_v2_session_factory(engine: Engine) -> sessionmaker[Session]: """Build a factory whose ordinary sessions are SQLModel Sessions.""" - return sessionmaker(bind=engine, class_=Session, expire_on_commit=False) def _discard_inherited_state(current_pid: int) -> None: """Drop pool state inherited through a process fork without touching parent FDs.""" - global _engine, _engine_pid, _engine_url_fingerprint, _session_factory if _engine is not None: @@ -86,7 +82,6 @@ def _discard_inherited_state(current_pid: int) -> None: def get_v2_engine(settings: V2DatabaseSettings | None = None) -> Engine: """Return the one lazy v2 runtime engine owned by this process.""" - global _engine, _engine_pid, _engine_url_fingerprint current_pid = os.getpid() @@ -114,7 +109,6 @@ def get_v2_session_factory( settings: V2DatabaseSettings | None = None, ) -> sessionmaker[Session]: """Return the process-owned SQLModel Session factory without opening a session.""" - global _session_factory engine = get_v2_engine(settings) @@ -126,7 +120,6 @@ def get_v2_session_factory( def close_v2_database() -> None: """Dispose this process's v2 pool and forget its lazy configuration.""" - global _engine, _engine_pid, _engine_url_fingerprint, _session_factory with _state_lock: diff --git a/policyengine_api/data/v2/household_migration_qualification.py b/policyengine_api/data/v2/household_migration_qualification.py index 11ca1616e..c680c8eca 100644 --- a/policyengine_api/data/v2/household_migration_qualification.py +++ b/policyengine_api/data/v2/household_migration_qualification.py @@ -142,14 +142,12 @@ def count(statement: object, table_name: str) -> int: def require_no_retained_household_data(counts: HouseholdDataCounts) -> None: """Stop the migration when any affected row requires a decision.""" - if counts.total: raise RetainedHouseholdDataError(counts) def build_qualification_engine(settings: V2DatabaseSettings) -> Engine: """Build an isolated connection pool for one qualification attempt.""" - return create_engine(settings.connection.url, poolclass=NullPool) @@ -165,7 +163,6 @@ def _read_and_require_counts(connection: Connection) -> HouseholdDataCounts: def household_migration_is_pending(connection: Connection) -> bool: """Return whether the Stage 11 household revision is in the upgrade path.""" - current_heads = MigrationContext.configure(connection).get_current_heads() revisions = ScriptDirectory.from_config( Config(V2_ALEMBIC_CONFIG) @@ -183,7 +180,6 @@ def qualify_household_migration_target( ), ) -> HouseholdMigrationQualification: """Qualify the configured Supabase target without changing it.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -212,7 +208,6 @@ def qualify_household_migration_if_pending( pending_checker: Callable[[Connection], bool] = household_migration_is_pending, ) -> HouseholdMigrationQualification: """Qualify affected rows only while the Stage 11 revision is pending.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -257,7 +252,6 @@ def _error_payload(error: Exception) -> dict[str, object]: def main() -> int: """Run qualification and return a shell-compatible status.""" - try: evidence = qualify_household_migration_if_pending() except Exception as error: # noqa: BLE001 - command must emit safe evidence diff --git a/policyengine_api/data/v2/metadata_validation.py b/policyengine_api/data/v2/metadata_validation.py index a37e92aaf..786908eb4 100644 --- a/policyengine_api/data/v2/metadata_validation.py +++ b/policyengine_api/data/v2/metadata_validation.py @@ -41,7 +41,6 @@ class V2MetadataValidationError(RuntimeError): def validate_v2_metadata_table_names(table_names: Iterable[str]) -> None: """Reject known v1-only and explicitly prohibited v2 table names.""" - actual = frozenset(table_names) prohibited = sorted(actual & PROHIBITED_V2_TABLES) v1_only = sorted(actual & V1_ONLY_TABLES) diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index 5595a3dee..bc3b9d646 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -88,7 +88,6 @@ def load_v2_alembic_settings( environ: Mapping[str, str] | None = None, ) -> V2AlembicSettings: """Select either the configured persistent target or isolated test Postgres.""" - values = os.environ if environ is None else environ raw_url = _required(values, V2_MIGRATION_DATABASE_URL) url = _parse_url(raw_url) @@ -134,7 +133,6 @@ def qualify_v2_connection( settings: V2AlembicSettings, ) -> None: """Verify live target identity and first-use state before Alembic runs.""" - if connection.dialect.name != "postgresql": raise V2MigrationTargetError("v2 Alembic requires a Postgres connection") if settings.disposable_test: @@ -172,7 +170,6 @@ def qualify_v2_connection( def validate_v2_head_schema(connection: Connection, metadata: MetaData) -> None: """Require live application tables to match authoritative ORM metadata.""" - public_tables = set(inspect(connection).get_table_names(schema="public")) application_tables = public_tables - {"alembic_version"} metadata_tables = {table.name for table in metadata.tables.values()} diff --git a/policyengine_api/data/v2/models/base.py b/policyengine_api/data/v2/models/base.py index 23c3809a9..fae3a25e2 100644 --- a/policyengine_api/data/v2/models/base.py +++ b/policyengine_api/data/v2/models/base.py @@ -45,7 +45,6 @@ def utc_now() -> datetime: """Return an aware UTC timestamp for application-created rows.""" - return datetime.now(UTC) @@ -78,7 +77,6 @@ class TimestampedModel(IdentifiedModel): def enum_type(enum_class: type, name: str) -> sa.Enum: """Build a stable lowercase-value Postgres enum for a string Enum.""" - # Native named enums and value serialization are SQLAlchemy features that # SQLModel intentionally exposes through Field(sa_type=...). return sa.Enum( diff --git a/policyengine_api/data/v2/policy_migration_qualification.py b/policyengine_api/data/v2/policy_migration_qualification.py index b88c237e7..5a514549e 100644 --- a/policyengine_api/data/v2/policy_migration_qualification.py +++ b/policyengine_api/data/v2/policy_migration_qualification.py @@ -125,14 +125,12 @@ def table_exists(name: str) -> bool: def require_no_retained_policy_data(counts: PolicyDataCounts) -> None: """Stop the migration when any predecessor policy row requires a decision.""" - if counts.total: raise RetainedPolicyDataError(counts) def build_qualification_engine(settings: V2DatabaseSettings) -> Engine: """Build an isolated connection pool for one qualification attempt.""" - return create_engine(settings.connection.url, poolclass=NullPool) @@ -157,7 +155,6 @@ def _qualify_connection(connection: Connection) -> PolicyDataCounts: def policy_migration_is_pending(connection: Connection) -> bool: """Return whether the Phase 10 policy revision is in the upgrade path.""" - current_heads = MigrationContext.configure(connection).get_current_heads() revisions = ScriptDirectory.from_config( Config(V2_ALEMBIC_CONFIG) @@ -173,7 +170,6 @@ def qualify_policy_migration_target( ), ) -> PolicyMigrationQualification: """Qualify the configured Supabase target without changing it.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -197,7 +193,6 @@ def qualify_policy_migration_if_pending( pending_checker: Callable[[Connection], bool] = policy_migration_is_pending, ) -> PolicyMigrationQualification: """Qualify predecessor rows only while the Phase 10 revision is pending.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -242,7 +237,6 @@ def _error_payload(error: Exception) -> dict[str, object]: def main() -> int: """Run qualification and return a shell-compatible status.""" - try: evidence = qualify_policy_migration_if_pending() except Exception as error: # noqa: BLE001 - command must emit safe evidence diff --git a/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py index e4da22634..6aa32f955 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -55,7 +55,6 @@ def set_report_type( report_type: str | None, ) -> Report: """Set a report's optional type only while it has no execution history.""" - report = _locked_report(session, report_id) if report.type == report_type: return report @@ -84,7 +83,6 @@ def create_report_run( idempotency_key: UUID | None = None, ) -> ReportRun: """Create one run, or return the run for a retried idempotent request.""" - _locked_report(session, report_id) if trigger is ReportRunTrigger.MANUAL and idempotency_key is None: raise ValueError("manual report reruns require an idempotency key") @@ -130,7 +128,6 @@ def begin_report_run( started_at: datetime | None = None, ) -> ReportRun: """Start or resume the same pending/running durable run after worker retry.""" - run = session.exec( select(ReportRun).where(ReportRun.id == report_run_id).with_for_update() ).one_or_none() @@ -154,7 +151,6 @@ def complete_report_run( markdown: str | None = None, ) -> ReportRun: """Mark the selected durable run successful without replacing history.""" - run = begin_report_run(session, report_run_id=report_run_id) run.status = ReportRunStatus.SUCCEEDED run.completed_at = completed_at or utc_now() @@ -173,7 +169,6 @@ def fail_report_run( completed_at: datetime | None = None, ) -> ReportRun: """Mark the selected durable run failed without deleting older success.""" - run = begin_report_run(session, report_run_id=report_run_id) run.status = ReportRunStatus.FAILED run.completed_at = completed_at or utc_now() @@ -191,7 +186,6 @@ def select_current_report_run( policyengine_version: str = POLICYENGINE_VERSION, ) -> ReportRun | None: """Select the deterministic current successful run for deployed versions.""" - report = session.get(Report, report_id) if report is None: raise ReportNotFoundError(f"report {report_id} does not exist") diff --git a/policyengine_api/data/v2/settings.py b/policyengine_api/data/v2/settings.py index 4ef27980f..98c852dea 100644 --- a/policyengine_api/data/v2/settings.py +++ b/policyengine_api/data/v2/settings.py @@ -49,13 +49,11 @@ class PostgresConnectionSettings: @property def url(self) -> URL: """Return SQLAlchemy's immutable URL for engine construction.""" - return self._url @property def redacted_url(self) -> str: """Return a log-safe URL with the password hidden.""" - return self._url.render_as_string(hide_password=True) def __str__(self) -> str: @@ -92,7 +90,6 @@ def _required(environ: Mapping[str, str], name: str) -> str: @cache def _load_secret_from_secret_manager(resource_name: str) -> str: """Resolve a v2 runtime URL only when a preview request selects it.""" - from google.cloud import secretmanager client = secretmanager.SecretManagerServiceClient() @@ -140,7 +137,6 @@ def load_supabase_target_settings( environ: Mapping[str, str] | None = None, ) -> SupabaseTargetSettings: """Load the externally configured non-secret Supabase target identity.""" - values = _environment(environ) project_ref = _required(values, V2_SUPABASE_PROJECT_REF) environment = _required(values, V2_SUPABASE_ENVIRONMENT) @@ -165,7 +161,6 @@ def parse_persistent_postgres_url( setting_name: str, ) -> PostgresConnectionSettings: """Validate an explicit persistent Postgres URL without echoing it.""" - try: url = make_url(raw_url) except ArgumentError as error: @@ -200,7 +195,6 @@ def validate_supabase_database_identity( setting_name: str, ) -> None: """Require a persistent URL to identify the configured Supabase project.""" - url = connection.url direct_host = f"db.{target.project_ref}.supabase.co" is_direct = url.host == direct_host @@ -245,7 +239,6 @@ def load_v2_runtime_database_settings( secret_loader: Callable[[str], str] | None = None, ) -> V2DatabaseSettings: """Load the future ordinary-runtime Postgres identity explicitly.""" - values = _environment(environ) raw_url = _resolve_runtime_database_url( values, @@ -269,7 +262,6 @@ def load_v2_migration_database_settings( environ: Mapping[str, str] | None = None, ) -> V2DatabaseSettings: """Load the schema-migration Postgres identity explicitly.""" - values = _environment(environ) return _database_settings( _required(values, V2_MIGRATION_DATABASE_URL), @@ -282,7 +274,6 @@ def load_v2_data_write_database_settings( environ: Mapping[str, str] | None = None, ) -> V2DatabaseSettings: """Load the one-time catalog row-write Postgres identity explicitly.""" - values = _environment(environ) return _database_settings( _required(values, V2_DATA_WRITE_DATABASE_URL), diff --git a/policyengine_api/dataset_display.py b/policyengine_api/dataset_display.py index b303ed3bd..d092fa7e8 100644 --- a/policyengine_api/dataset_display.py +++ b/policyengine_api/dataset_display.py @@ -11,7 +11,6 @@ def get_dataset_display_label(dataset_name: object) -> str: """Return the public label for an internal logical dataset name.""" - normalized_name = str(dataset_name or "") for prefix, label in DATASET_FAMILY_DISPLAY_LABELS: if normalized_name.startswith(prefix): diff --git a/policyengine_api/fastapi_routes/responses.py b/policyengine_api/fastapi_routes/responses.py index a58ab27ec..18ab9e249 100644 --- a/policyengine_api/fastapi_routes/responses.py +++ b/policyengine_api/fastapi_routes/responses.py @@ -13,7 +13,6 @@ def render_legacy_json(content: JSONValue) -> bytes: """Serialize once with the same JSON encoder used by legacy metadata.""" - return json.dumps(content).encode("utf-8") diff --git a/policyengine_api/fastapi_routes/v2/errors.py b/policyengine_api/fastapi_routes/v2/errors.py index 4a357c961..65bbf80b4 100644 --- a/policyengine_api/fastapi_routes/v2/errors.py +++ b/policyengine_api/fastapi_routes/v2/errors.py @@ -21,7 +21,6 @@ class V2ErrorResponse(BaseModel): def v2_error_response(status_code: int, message: str) -> JSONResponse: """Serialize one API v2 error without resource-path dispatch.""" - error = V2ErrorResponse(message=message) return JSONResponse( status_code=status_code, diff --git a/policyengine_api/fastapi_routes/v2/households/routes.py b/policyengine_api/fastapi_routes/v2/households/routes.py index cf0fe4006..6e17e322a 100644 --- a/policyengine_api/fastapi_routes/v2/households/routes.py +++ b/policyengine_api/fastapi_routes/v2/households/routes.py @@ -50,7 +50,6 @@ async def enforce_household_request_size(request: Request) -> None: """Reject declared or actual request bodies larger than 1 MiB.""" - content_length = request.headers.get("content-length") if content_length is not None: try: @@ -103,7 +102,6 @@ def build_v2_household_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build native household routes without opening a database connection.""" - router = APIRouter(prefix="/v2", responses=HOUSEHOLD_ERROR_RESPONSES) create_query = query_dependency(HouseholdCreateQuery) detail_query = query_dependency(HouseholdDetailQuery) diff --git a/policyengine_api/fastapi_routes/v2/policies/routes.py b/policyengine_api/fastapi_routes/v2/policies/routes.py index 3ecd854ec..9b33d467a 100644 --- a/policyengine_api/fastapi_routes/v2/policies/routes.py +++ b/policyengine_api/fastapi_routes/v2/policies/routes.py @@ -55,7 +55,6 @@ async def enforce_policy_request_size(request: Request) -> None: """Bound both declared and actual request bytes before service creation.""" - content_length = request.headers.get("content-length") if content_length is not None: try: @@ -108,7 +107,6 @@ def build_v2_policy_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build native policy routes without opening a database connection.""" - router = APIRouter(prefix="/v2", responses=POLICY_ERROR_RESPONSES) create_query = query_dependency(PolicyCreateQuery) detail_query = query_dependency(PolicyDetailQuery) diff --git a/policyengine_api/fastapi_routes/v2/routes.py b/policyengine_api/fastapi_routes/v2/routes.py index 85f4dbff5..9698ad5d0 100644 --- a/policyengine_api/fastapi_routes/v2/routes.py +++ b/policyengine_api/fastapi_routes/v2/routes.py @@ -36,7 +36,6 @@ def build_v2_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build isolated resource routes without loading v2 configuration.""" - router = APIRouter() router.include_router(build_v2_user_household_router(dependencies)) router.include_router(build_v2_user_policy_router(dependencies)) diff --git a/policyengine_api/fastapi_routes/v2/user_households/routes.py b/policyengine_api/fastapi_routes/v2/user_households/routes.py index 8710ac1f7..2d42a9246 100644 --- a/policyengine_api/fastapi_routes/v2/user_households/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_households/routes.py @@ -89,7 +89,6 @@ def build_v2_user_household_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build native association routes without opening a database connection.""" - router = APIRouter(prefix="/v2", responses=USER_HOUSEHOLD_ERROR_RESPONSES) create_query = query_dependency(UserHouseholdCreateQuery) detail_query = query_dependency(UserHouseholdDetailQuery) diff --git a/policyengine_api/fastapi_routes/v2/user_policies/routes.py b/policyengine_api/fastapi_routes/v2/user_policies/routes.py index 740eeedaf..c1ae65c03 100644 --- a/policyengine_api/fastapi_routes/v2/user_policies/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_policies/routes.py @@ -85,7 +85,6 @@ def build_v2_user_policy_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build native association routes without opening a database connection.""" - router = APIRouter(prefix="/v2", responses=USER_POLICY_ERROR_RESPONSES) country_query = query_dependency(CountryQuery) collection_query = query_dependency(UserPolicyCollectionQuery) diff --git a/policyengine_api/gcp_logging.py b/policyengine_api/gcp_logging.py index 6c3b55b7e..8e89aafe2 100644 --- a/policyengine_api/gcp_logging.py +++ b/policyengine_api/gcp_logging.py @@ -25,7 +25,6 @@ def log_struct( labels: Mapping[str, Any] | None = None, ) -> None: """Record an allowlisted structured message without affecting callers.""" - try: message = str(info.get("message") or "API event") attributes = { diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 9e33a8cfe..1eead55d5 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -60,7 +60,6 @@ def resolve_simulation_entrypoint_url(entrypoint: str | None = None) -> str: def _attach_current_request_id(request: httpx.Request) -> None: """Attach the active API request ID to an outbound simulation request.""" - request_id = current_request_id() if request_id is not None: request.headers[REQUEST_ID_HEADER] = request_id @@ -71,9 +70,7 @@ def _attach_current_request_id(request: httpx.Request) -> None: @dataclass class ModalSimulationExecution: - """ - Represents a Modal simulation job execution. - """ + """Represents a Modal simulation job execution.""" job_id: str status: str @@ -90,9 +87,7 @@ def name(self) -> str: @dataclass class ModalBudgetWindowBatchExecution: - """ - Represents a budget-window batch execution in the Simulation Entrypoint. - """ + """Represents a budget-window batch execution in the Simulation Entrypoint.""" batch_job_id: str status: str @@ -111,8 +106,7 @@ def name(self) -> str: class SimulationEntrypointClient: - """ - HTTP client for the configured Simulation Entrypoint. + """HTTP client for the configured Simulation Entrypoint. This class provides methods for submitting simulation jobs and polling for their status/results via HTTP endpoints. @@ -191,8 +185,7 @@ def _normalize_submission_payload(self, payload: dict) -> dict: return modal_payload def run(self, payload: dict) -> ModalSimulationExecution: - """ - Submit a simulation job to the selected Simulation entrypoint. + """Submit a simulation job to the selected Simulation entrypoint. Parameters ---------- @@ -209,6 +202,7 @@ def run(self, payload: dict) -> ModalSimulationExecution: ------ httpx.HTTPStatusError If the API returns an error response. + """ try: modal_payload = self._normalize_submission_payload(payload) @@ -260,9 +254,7 @@ def run(self, payload: dict) -> ModalSimulationExecution: raise def run_budget_window_batch(self, payload: dict) -> ModalBudgetWindowBatchExecution: - """ - Submit a budget-window batch job to the selected Simulation entrypoint. - """ + """Submit a budget-window batch job to the selected Simulation entrypoint.""" try: modal_payload = self._normalize_submission_payload(payload) @@ -404,8 +396,7 @@ def resolve_app_name( ) from exc def get_execution_id(self, execution: ModalSimulationExecution) -> str: - """ - Get the job ID from an execution. + """Get the job ID from an execution. Parameters ---------- @@ -416,12 +407,12 @@ def get_execution_id(self, execution: ModalSimulationExecution) -> str: ------- str The job ID. + """ return execution.job_id def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: - """ - Poll the selected Simulation entrypoint for the current status of a job. + """Poll the selected Simulation entrypoint for the current status of a job. Parameters ---------- @@ -432,6 +423,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: ------- ModalSimulationExecution Execution object with current status and result if complete. + """ try: response = self.client.get(f"{self.base_url}/jobs/{job_id}") @@ -471,9 +463,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: def get_budget_window_batch_by_id( self, batch_job_id: str ) -> ModalBudgetWindowBatchExecution: - """ - Poll the selected Simulation entrypoint for a budget-window batch. - """ + """Poll the selected Simulation entrypoint for a budget-window batch.""" try: response = self.client.get( f"{self.base_url}/budget-window-jobs/{batch_job_id}" @@ -515,8 +505,7 @@ def get_budget_window_batch_by_id( raise def get_execution_status(self, execution: ModalSimulationExecution) -> str: - """ - Get the status string from an execution. + """Get the status string from an execution. Parameters ---------- @@ -527,12 +516,12 @@ def get_execution_status(self, execution: ModalSimulationExecution) -> str: ------- str The status string ("submitted", "running", "complete", "failed"). + """ return execution.status def get_execution_result(self, execution: ModalSimulationExecution) -> dict | None: - """ - Get the result from a completed execution. + """Get the result from a completed execution. Parameters ---------- @@ -543,17 +532,18 @@ def get_execution_result(self, execution: ModalSimulationExecution) -> dict | No ------- dict or None The simulation result if complete, None otherwise. + """ return execution.result def health_check(self) -> bool: - """ - Check if the selected Simulation entrypoint is healthy. + """Check if the selected Simulation entrypoint is healthy. Returns ------- bool True if the API is healthy, False otherwise. + """ try: response = self.client.get(f"{self.base_url}/health") diff --git a/policyengine_api/migration_flags.py b/policyengine_api/migration_flags.py index 34843a1c2..c1d5cbab3 100644 --- a/policyengine_api/migration_flags.py +++ b/policyengine_api/migration_flags.py @@ -149,7 +149,6 @@ def get_db_read(entity: str) -> str: def get_v1_policy_write_source() -> str: """Select Cloud SQL alone or immediate Cloud SQL-to-Supabase mirroring.""" - return _read_choice( "DB_WRITE_POLICY", DEFAULT_DB_SOURCE, @@ -159,7 +158,6 @@ def get_v1_policy_write_source() -> str: def get_v1_policy_read_source() -> str: """Require every v1 policy read to remain on Cloud SQL in Phase 10.""" - return _read_choice( "DB_READ_POLICY", DEFAULT_DB_SOURCE, @@ -169,7 +167,6 @@ def get_v1_policy_read_source() -> str: def get_v1_household_write_source() -> str: """Select Cloud SQL alone or immediate copying to Supabase.""" - return _read_choice( "DB_WRITE_HOUSEHOLD", DEFAULT_DB_SOURCE, @@ -179,7 +176,6 @@ def get_v1_household_write_source() -> str: def get_v1_household_read_source() -> str: """Require every v1 household read to remain on Cloud SQL in Phase 11.""" - return _read_choice( "DB_READ_HOUSEHOLD", DEFAULT_DB_SOURCE, diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py index c521e77a8..ce8a25524 100644 --- a/policyengine_api/migration_logging.py +++ b/policyengine_api/migration_logging.py @@ -156,7 +156,6 @@ def log_migration_request( route_impl: RouteImplementation | None = None, ) -> None: """Log a migration-aware API request in the shared structured format.""" - elapsed_ms = None if started_at is not None: elapsed_ms = round((time.time() - started_at) * 1000, 2) diff --git a/policyengine_api/observability/identifiers.py b/policyengine_api/observability/identifiers.py index 57be0c593..64d8899d0 100644 --- a/policyengine_api/observability/identifiers.py +++ b/policyengine_api/observability/identifiers.py @@ -9,13 +9,11 @@ def generate_observability_id() -> str: """Create an identifier used only to correlate observability records.""" - return str(uuid4()) def normalize_observability_id(value: object) -> str | None: """Return a canonical UUID string, or ``None`` for malformed input.""" - if not isinstance(value, str): return None try: diff --git a/policyengine_api/observability/runtime.py b/policyengine_api/observability/runtime.py index aea6f10c2..b944ab3eb 100644 --- a/policyengine_api/observability/runtime.py +++ b/policyengine_api/observability/runtime.py @@ -73,7 +73,6 @@ def get_runtime() -> ObservabilityRuntime: def set_runtime_context(**attributes: object) -> None: """Bind local telemetry attributes without affecting application behavior.""" - try: runtime.set_context(**attributes) except Exception: diff --git a/policyengine_api/query_parameters.py b/policyengine_api/query_parameters.py index c11946886..81c5c83e4 100644 --- a/policyengine_api/query_parameters.py +++ b/policyengine_api/query_parameters.py @@ -45,13 +45,11 @@ def normalize_country_id(value: object) -> object: """Lowercase a textual country ID before its supported-value check.""" - return value.lower() if isinstance(value, str) else value def validate_legacy_user_id(value: str) -> str: """Reject an empty or whitespace-only legacy user identifier.""" - if not value.strip(): raise ValueError( "legacy user_id must contain at least one non-whitespace character" @@ -285,7 +283,6 @@ def query_field_multiplicity( model_type: type[StrictQueryParameters], ) -> Mapping[str, bool]: """Return public query names mapped to whether repeated values are valid.""" - return { (field.alias or field_name): _annotation_is_list(field.annotation) for field_name, field in model_type.model_fields.items() @@ -297,7 +294,6 @@ def validate_scalar_query_multiplicity( items: Iterable[tuple[str, str]], ) -> None: """Reject a repeated declared scalar while allowing declared list fields.""" - multiplicity = query_field_multiplicity(model_type) seen: set[str] = set() for name, _value in items: @@ -313,7 +309,6 @@ def parse_query_items( items: Iterable[tuple[str, str]], ) -> QueryParametersT: """Validate ordered query pairs through one canonical Pydantic schema.""" - pairs = list(items) validate_scalar_query_multiplicity(model_type, pairs) multiplicity = query_field_multiplicity(model_type) @@ -334,7 +329,6 @@ def parse_multidict_query( query_input: object, ) -> QueryParametersT: """Adapt a Flask/Werkzeug-style MultiDict to the canonical parser.""" - items_method = getattr(query_input, "items", None) if not callable(items_method): raise TypeError("query input must provide an items method") diff --git a/policyengine_api/readiness.py b/policyengine_api/readiness.py index dfcdf2325..77daf6dbd 100644 --- a/policyengine_api/readiness.py +++ b/policyengine_api/readiness.py @@ -16,7 +16,6 @@ def validate_policy_runtime_configuration() -> None: """Validate v1 source selectors, v2 persistence and measurement settings.""" - from policyengine_api.data.v2.settings import ( load_v2_runtime_database_settings, ) diff --git a/policyengine_api/request_context.py b/policyengine_api/request_context.py index 6e7fd8ef3..8c344594d 100644 --- a/policyengine_api/request_context.py +++ b/policyengine_api/request_context.py @@ -30,13 +30,11 @@ def generate_request_id() -> str: """Generate an API request correlation identifier.""" - return uuid.uuid4().hex def current_request_id() -> str | None: """Return the correlation identifier for the current request, if any.""" - if flask.has_request_context(): return getattr(flask.g, "request_id", None) return _asgi_request_id.get() @@ -44,7 +42,6 @@ def current_request_id() -> str | None: def current_observability_id() -> str | None: """Return the diagnostic correlation identifier for the current request.""" - if flask.has_request_context(): return getattr(flask.g, "observability_id", None) return _asgi_observability_id.get() @@ -52,7 +49,6 @@ def current_observability_id() -> str | None: def incoming_observability_id() -> str | None: """Return the validated identifier candidate supplied by the caller.""" - if flask.has_request_context(): return getattr(flask.g, "incoming_observability_id", None) return _asgi_incoming_observability_id.get() @@ -74,7 +70,6 @@ def _bind_observability_id(observability_id: str) -> str: def start_observability_id(value: object = None) -> str: """Bind the identifier for a newly accepted calculation or report.""" - current = current_observability_id() if current is not None: return current @@ -86,7 +81,6 @@ def start_observability_id(value: object = None) -> str: def restore_observability_id(value: object) -> str | None: """Bind a valid identifier read from durable functional state.""" - observability_id = normalize_observability_id(value) if observability_id is None: return current_observability_id() @@ -95,5 +89,4 @@ def restore_observability_id(value: object) -> str | None: def resolve_observability_id(value: object) -> str: """Use a valid caller value or create a new diagnostic identifier.""" - return normalize_observability_id(value) or generate_observability_id() diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index 7eefd896e..4d7ff2f13 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -258,12 +258,12 @@ def _calculation_cache_key(*args: object, **kwargs: object) -> str: @household_bp.route("//household/", methods=["GET"]) @validate_country def get_household(country_id: str, household_id: int) -> Response: - """ - Get a household's input data with a given ID. + """Get a household's input data with a given ID. Args: country_id (str): The country ID. household_id (int): The household ID. + """ try: get_v1_household_read_source() @@ -290,13 +290,12 @@ def get_household(country_id: str, household_id: int) -> Response: @household_bp.route("//household", methods=["POST"]) @validate_country def post_household(country_id: str) -> Response: - """ - Set a household's input data. + """Set a household's input data. Args: country_id (str): The country ID. - """ + """ # Validate payload payload = request.json is_payload_valid, message = validate_household_payload(payload) diff --git a/policyengine_api/routes/metadata_routes.py b/policyengine_api/routes/metadata_routes.py index 274585fe4..f70e37bfd 100644 --- a/policyengine_api/routes/metadata_routes.py +++ b/policyengine_api/routes/metadata_routes.py @@ -20,8 +20,8 @@ def get_metadata(country_id: str) -> Response: Args: country_id (str): The country ID. - """ + """ # Retrieve country metadata and add status and message to the response country_metadata = metadata_service.get_metadata(country_id) body = json.dumps( diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index d27884f2a..461da8c0e 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -41,7 +41,6 @@ def _should_mirror_to_v2(country_id: str, write_source: str) -> bool: """Select immediate mirroring only for countries supported by API v2.""" - return write_source == "dual_write" and country_id in SUPPORTED_V2_COUNTRY_IDS @@ -87,7 +86,6 @@ def _user_policy_persistence_failure( legacy_user_policy_id: object | None = None, ) -> Response: """Return and record a failure without serializing exception details.""" - status_code = 503 if error.retryable else 500 try: logger.log_struct( @@ -144,8 +142,7 @@ def _serialize_policy(policy: Policy) -> dict: @policy_bp.route("//policy/", methods=["GET"]) @validate_country def get_policy(country_id: str, policy_id: int | str) -> Response: - """ - Get policy data for a given country and policy ID. + """Get policy data for a given country and policy ID. Args: country_id (str) @@ -154,8 +151,8 @@ def get_policy(country_id: str, policy_id: int | str) -> Response: Returns: Response: A Flask response object containing the policy data in JSON format - """ + """ try: get_v1_policy_read_source() except ValueError: @@ -179,14 +176,13 @@ def get_policy(country_id: str, policy_id: int | str) -> Response: @policy_bp.route("//policy", methods=["POST"]) @validate_country def set_policy(country_id: str) -> Response: - """ - Set policy data for given country and policy. If policy already exists, + """Set policy data for given country and policy. If policy already exists, return existing policy and 200. Args: country_id (str) - """ + """ payload = request.json is_payload_valid, message = validate_set_policy_payload(payload) diff --git a/policyengine_api/routes/reform_impact_routes.py b/policyengine_api/routes/reform_impact_routes.py index c3a9d2921..82b2a9f46 100644 --- a/policyengine_api/routes/reform_impact_routes.py +++ b/policyengine_api/routes/reform_impact_routes.py @@ -24,7 +24,6 @@ def _serialize_v1_datetime(value: datetime | None) -> str | None: def _serialize_v1_reform_impact(impact: CachedReformImpact) -> dict[str, object]: """Project a cached impact onto the explicit historical v1 response shape.""" - return { "reform_impact_id": impact.reform_impact_id, "baseline_policy_id": impact.baseline_policy_id, diff --git a/policyengine_api/routes/report_output_routes.py b/policyengine_api/routes/report_output_routes.py index b105b33ba..ea188dcc4 100644 --- a/policyengine_api/routes/report_output_routes.py +++ b/policyengine_api/routes/report_output_routes.py @@ -19,7 +19,6 @@ def _serialize_v1_report_output(view: ReportOutputView) -> dict: """Project mapped report state onto the historical v1 response shape.""" - report_output = view.report_output result = { column.name: getattr(report_output, column.name) @@ -40,8 +39,7 @@ def _serialize_v1_report_output(view: ReportOutputView) -> dict: @report_output_bp.route("//report", methods=["POST"]) @validate_country def create_report_output(country_id: str) -> Response: - """ - Create a new report output record with pending status. + """Create a new report output record with pending status. Args: country_id (str): The country ID. @@ -50,6 +48,7 @@ def create_report_output(country_id: str) -> Response: - simulation_1_id (int): The first simulation ID (required) - simulation_2_id (int | None): The second simulation ID (optional, for comparisons) - year (str | None): The year for the report (optional) + """ print(f"Creating report output for country {country_id}") payload = request.json @@ -117,8 +116,7 @@ def create_report_output(country_id: str) -> Response: @report_output_bp.route("//report/", methods=["GET"]) @validate_country def get_report_output(country_id: str, report_id: int) -> Response: - """ - Get a report output record by ID. + """Get a report output record by ID. The response result may include requested_at, started_at, and finished_at values projected from the selected report_output_runs row. Those fields are @@ -128,6 +126,7 @@ def get_report_output(country_id: str, report_id: int) -> Response: Args: country_id (str): The country ID. report_id (int): The report output ID. + """ print(f"Getting report output {report_id} for country {country_id}") @@ -152,8 +151,7 @@ def get_report_output(country_id: str, report_id: int) -> Response: @report_output_bp.route("//report", methods=["PATCH"]) @validate_country def update_report_output(country_id: str) -> Response: - """ - Update a report output record with results or error. + """Update a report output record with results or error. Args: country_id (str): The country ID. @@ -164,8 +162,8 @@ def update_report_output(country_id: str) -> Response: - output (dict): The result output (for complete status) - api_version (str): The API version of the report - error_message (str): The error message (for error status) - """ + """ payload = request.json if payload is None: raise BadRequest("Payload missing from request") diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index 5e9b329e6..f0853c53b 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -28,7 +28,6 @@ def _spm_override_response(payload: dict) -> Response | None: def _serialize_v1_simulation(simulation: Simulation) -> dict: """Project canonical ORM JSON objects onto the legacy v1 response shape.""" - response = { column.name: getattr(simulation, column.name) for column in Simulation.__table__.columns @@ -47,8 +46,7 @@ def _serialize_v1_simulation(simulation: Simulation) -> dict: @simulation_bp.route("//simulation", methods=["POST"]) @validate_country def create_simulation(country_id: str) -> Response: - """ - Create a new simulation record. + """Create a new simulation record. Args: country_id (str): The country ID. @@ -57,6 +55,7 @@ def create_simulation(country_id: str) -> Response: - population_id (str): The population identifier (household or geography ID) - population_type (str): Type of population ('household' or 'geography') - policy_id (int): The policy ID + """ print(f"Creating simulation for country {country_id}") @@ -129,12 +128,12 @@ def create_simulation(country_id: str) -> Response: @simulation_bp.route("//simulation/", methods=["GET"]) @validate_country def get_simulation(country_id: str, simulation_id: int) -> Response: - """ - Get a simulation record by ID. + """Get a simulation record by ID. Args: country_id (str): The country ID. simulation_id (int): The simulation ID. + """ print(f"Getting simulation {simulation_id} for country {country_id}") @@ -165,8 +164,7 @@ def get_simulation(country_id: str, simulation_id: int) -> Response: @simulation_bp.route("//simulation", methods=["PATCH"]) @validate_country def update_simulation(country_id: str) -> Response: - """ - Update a simulation record with results or error. + """Update a simulation record with results or error. Args: country_id (str): The country ID. @@ -177,8 +175,8 @@ def update_simulation(country_id: str) -> Response: - output (dict): The result output (for complete status) - api_version (str): The API version of the simulation - error_message (str): The error message (for error status) - """ + """ payload = request.json if payload is None: raise BadRequest("Payload missing from request") diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index 8e36f9707..8bcd68bd7 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -30,10 +30,7 @@ def _serialize_user_profile( @user_profile_bp.route("//user-profile", methods=["POST"]) @validate_country def set_user_profile(country_id: str) -> Response: - """ - Creates a new user_profile - """ - + """Creates a new user_profile""" payload = request.json if payload is None: raise BadRequest("Payload missing from request") @@ -104,12 +101,10 @@ def get_user_profile(country_id: str) -> Response: @user_profile_bp.route("//user-profile", methods=["PUT"]) @validate_country def update_user_profile(country_id: str) -> Response: - """ - Update any part of a user_profile, given a user_id, + """Update any part of a user_profile, given a user_id, except the auth0_id value; any attempt to edit this will assume malicious intent and 403 """ - payload = request.json if payload is None: diff --git a/policyengine_api/runtime_cache/client.py b/policyengine_api/runtime_cache/client.py index e6bbf0dae..a7225db18 100644 --- a/policyengine_api/runtime_cache/client.py +++ b/policyengine_api/runtime_cache/client.py @@ -40,7 +40,6 @@ def get_runtime_cache_client( settings: RuntimeCacheSettings | None = None, ) -> redis.Redis: """Return one client per process without connecting during import.""" - process_id = os.getpid() if process_id not in _clients: with _client_lock: diff --git a/policyengine_api/runtime_cache/core.py b/policyengine_api/runtime_cache/core.py index 7d50bdfda..058c31b9d 100644 --- a/policyengine_api/runtime_cache/core.py +++ b/policyengine_api/runtime_cache/core.py @@ -127,7 +127,6 @@ def jittered_ttl( Coordination and claim TTLs are safety bounds and must not use this helper. """ - if ttl_seconds <= 0: raise ValueError("cache TTL must be positive") if not 0 <= jitter_fraction < 1: @@ -148,7 +147,6 @@ def record_cache_event( operation: str | None = None, ) -> None: """Emit one metric-ready event without accepting cache keys or values.""" - payload: dict[str, str | int | float] = { "message": "Runtime cache operation", "metric_name": "runtime_cache_operations", diff --git a/policyengine_api/runtime_cache/reform_impacts.py b/policyengine_api/runtime_cache/reform_impacts.py index e21bdfe2b..9012010b3 100644 --- a/policyengine_api/runtime_cache/reform_impacts.py +++ b/policyengine_api/runtime_cache/reform_impacts.py @@ -219,7 +219,6 @@ def claim_start( observability_id: str, ) -> bool: """Atomically claim ownership of one reform-impact submission.""" - claim = ReformImpactStartClaim.from_payload( { "submission_claim_id": claim_token, @@ -260,7 +259,6 @@ def get_start_claim( target: str, ) -> ReformImpactStartClaim | None: """Read the current submission owner and its diagnostic identifier.""" - started_at = time.perf_counter() key = self._start_claim_key( country_id=country_id, @@ -329,7 +327,6 @@ def release_start( observability_id: str, ) -> bool: """Release a start claim only when its complete state still matches.""" - claim = ReformImpactStartClaim.from_payload( { "submission_claim_id": claim_token, @@ -584,6 +581,5 @@ def delete_matching_computing( def reform_impact_id(execution_id: str) -> int: """Stable positive cache-local identifier for the historical response field.""" - digest = hashlib.sha256(execution_id.encode("utf-8")).hexdigest()[:15] return int(digest, 16) diff --git a/policyengine_api/runtime_cache/settings.py b/policyengine_api/runtime_cache/settings.py index 6fb2937da..d7b293d0f 100644 --- a/policyengine_api/runtime_cache/settings.py +++ b/policyengine_api/runtime_cache/settings.py @@ -129,7 +129,6 @@ def load_runtime_cache_settings( environ: Mapping[str, str] | None = None, ) -> RuntimeCacheSettings: """Load disabled, explicit-local, or fail-closed deployed cache settings.""" - values = os.environ if environ is None else environ raw_mode = values.get(RUNTIME_CACHE_MODE, "").strip() deployed_runtime = _is_deployed(values) diff --git a/policyengine_api/services/budget_window_cache.py b/policyengine_api/services/budget_window_cache.py index b6190341a..1ec50c10d 100644 --- a/policyengine_api/services/budget_window_cache.py +++ b/policyengine_api/services/budget_window_cache.py @@ -146,7 +146,6 @@ def _handle_cache_error( def get_state(self, cache_key: str) -> BudgetWindowCacheState | None: """Read the complete report state or fail closed on cache outage.""" - started_at = time.perf_counter() state_key = self._state_key(cache_key) try: diff --git a/policyengine_api/services/economy_comparison.py b/policyengine_api/services/economy_comparison.py index 08de03777..2d5518827 100644 --- a/policyengine_api/services/economy_comparison.py +++ b/policyengine_api/services/economy_comparison.py @@ -116,8 +116,7 @@ def detailed_budgetary_impact(baseline: dict, reform: dict, country_id: str) -> def decile_impact(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on the deciles of the population. + """Compare the impact of a reform on the deciles of the population. Args: baseline (dict): The baseline economy. @@ -125,6 +124,7 @@ def decile_impact(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on the deciles of the population. + """ baseline_income = MicroSeries( baseline["household_net_income"], weights=baseline["household_weight"] @@ -157,8 +157,7 @@ def decile_impact(baseline: dict, reform: dict) -> dict: def wealth_decile_impact(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on the deciles of the population. + """Compare the impact of a reform on the deciles of the population. Args: baseline (dict): The baseline economy. @@ -166,6 +165,7 @@ def wealth_decile_impact(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on the deciles of the population. + """ baseline_income = MicroSeries( baseline["household_net_income"], weights=baseline["household_weight"] @@ -197,8 +197,7 @@ def wealth_decile_impact(baseline: dict, reform: dict) -> dict: def inequality_impact(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on inequality. + """Compare the impact of a reform on inequality. Args: baseline (dict): The baseline economy. @@ -206,8 +205,8 @@ def inequality_impact(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on inequality. - """ + """ return dict( gini=dict( baseline=baseline["gini"], @@ -225,8 +224,7 @@ def inequality_impact(baseline: dict, reform: dict) -> dict: def poverty_impact(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on poverty. + """Compare the impact of a reform on poverty. Args: baseline (dict): The baseline economy. @@ -234,6 +232,7 @@ def poverty_impact(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on poverty. + """ baseline_poverty = MicroSeries( baseline["person_in_poverty"], weights=baseline["person_weight"] @@ -297,7 +296,8 @@ def compute_income_change( baseline_values: np.ndarray, reform_values: np.ndarray ) -> np.ndarray: """Percentage income change with a floor of 1 on the baseline - to avoid division by zero for zero/negative incomes.""" + to avoid division by zero for zero/negative incomes. + """ absolute_change = reform_values - baseline_values capped_baseline = np.maximum(baseline_values, 1) return absolute_change / capped_baseline @@ -408,8 +408,7 @@ def intra_wealth_decile_impact(baseline: dict, reform: dict) -> dict: def poverty_gender_breakdown(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on poverty. + """Compare the impact of a reform on poverty. Args: baseline (dict): The baseline economy. @@ -417,6 +416,7 @@ def poverty_gender_breakdown(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on poverty. + """ if baseline["is_male"] is None: return {} @@ -463,8 +463,7 @@ def poverty_gender_breakdown(baseline: dict, reform: dict) -> dict: def poverty_racial_breakdown(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on poverty. + """Compare the impact of a reform on poverty. Args: baseline (dict): The baseline economy. @@ -472,6 +471,7 @@ def poverty_racial_breakdown(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on poverty. + """ if baseline["race"] is None: return {} @@ -785,8 +785,7 @@ def compare_economic_outputs( country_id: str = None, region: str | None = None, ) -> dict: - """ - Compare the economic outputs of two economies. + """Compare the economic outputs of two economies. Args: baseline (dict): The baseline economy. @@ -797,6 +796,7 @@ def compare_economic_outputs( Returns: dict: The comparison of the two economies. + """ if baseline.get("type") == "general": budgetary_impact_data = budgetary_impact(baseline, reform) diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 4c72feac6..fb7e81ecc 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -71,9 +71,7 @@ class ImpactAction(Enum): - """ - Enum for the action to take based on the status of an economic impact calculation. - """ + """Enum for the action to take based on the status of an economic impact calculation.""" COMPLETED = "completed" FAILED = "failed" @@ -82,9 +80,7 @@ class ImpactAction(Enum): class ImpactStatus(Enum): - """ - Enum for the status of an economic impact calculation. - """ + """Enum for the status of an economic impact calculation.""" COMPUTING = "computing" OK = "ok" @@ -134,8 +130,7 @@ class EconomicImpactSetupOptions(BaseModel): class EconomicImpactResult(BaseModel): - """ - Model for the result of an economic impact calculation. + """Model for the result of an economic impact calculation. Contains the status and the reform impact JSON. """ @@ -146,9 +141,7 @@ class EconomicImpactResult(BaseModel): model_config = {"frozen": True} # Make model immutable def to_dict(self) -> dict[str, str | dict | None]: - """ - Convert the EconomicImpactResult to a dictionary. - """ + """Convert the EconomicImpactResult to a dictionary.""" return { "status": self.status.value, "data": self.data, @@ -157,31 +150,23 @@ def to_dict(self) -> dict[str, str | dict | None]: @classmethod def completed(cls, data: dict) -> "EconomicImpactResult": - """ - Create an EconomicImpactResult for a completed impact calculation. - """ + """Create an EconomicImpactResult for a completed impact calculation.""" return cls(status=ImpactStatus.OK, data=data) @classmethod def computing(cls) -> "EconomicImpactResult": - """ - Create an EconomicImpactResult for a computing impact calculation. - """ + """Create an EconomicImpactResult for a computing impact calculation.""" return cls(status=ImpactStatus.COMPUTING, data=None) @classmethod def error(cls, message: str) -> "EconomicImpactResult": - """ - Create an EconomicImpactResult for an error in the impact calculation. - """ + """Create an EconomicImpactResult for an error in the impact calculation.""" logger.log_struct({"message": message}, severity="ERROR") return cls(status=ImpactStatus.ERROR, data=None, message=message) class BudgetWindowEconomicImpactResult(BaseModel): - """ - Model for a batch budget-window economic impact response. - """ + """Model for a batch budget-window economic impact response.""" status: ImpactStatus data: dict | None = None @@ -264,8 +249,7 @@ def failed( class EconomyService: - """ - Service for calculating economic impact of policy reforms; this is connected + """Service for calculating economic impact of policy reforms; this is connected to the /economy route, which does not have its own table; therefore, it connects with other services to access their respective tables """ @@ -328,7 +312,6 @@ def _get_policy_jsons( @staticmethod def _parse_json_object(value: JSONValue) -> dict[str, Any]: """Accept ORM-decoded objects and legacy JSON text at a read boundary.""" - parsed = json.loads(value) if isinstance(value, str) else value if not isinstance(parsed, dict): raise TypeError("Expected a JSON object") @@ -347,15 +330,13 @@ def get_economic_impact( api_version: str, target: Literal["general", "cliff"] = "general", ) -> EconomicImpactResult: - """ - Calculate the society-wide economic impact of a policy reform. + """Calculate the society-wide economic impact of a policy reform. Returns a tuple of: - status: "ok", "error", or "computing" - reform_impact_json: The JSON object containing the economic impact data, or None if the status is "computing" or "error". """ - set_runtime_context( country_id=country_id, policy_id=policy_id, @@ -1161,10 +1142,7 @@ def _get_previous_impacts( options_hash: str, api_version: str, ) -> list[CachedReformImpact]: - """ - Fetch any previous simulation runs for the given policy reform. - """ - + """Fetch any previous simulation runs for the given policy reform.""" return self._reform_impacts.get_all_reform_impacts_by_options_hash_prefix( country_id, policy_id, @@ -1181,8 +1159,7 @@ def _get_most_recent_impact( self, setup_options: EconomicImpactSetupOptions, ) -> CachedReformImpact | None: - """ - Get the first impact from the reform_impacts table based on the provided setup options. + """Get the first impact from the reform_impacts table based on the provided setup options. Returns the first impact if it exists, otherwise returns None. """ previous_impacts = self._get_previous_impacts( @@ -1231,8 +1208,7 @@ def _handle_execution_state( reform_impact: ReformImpact, execution: ModalSimulationExecution | None = None, ) -> EconomicImpactResult: - """ - Handle the state of the execution and return the appropriate status and result. + """Handle the state of the execution and return the appropriate status and result. Supports both GCP Workflow statuses (SUCCEEDED, FAILED, ACTIVE) and Modal statuses (complete, failed, running, submitted). @@ -1501,10 +1477,7 @@ def _setup_sim_options( policyengine_version: str | None = None, spm: dict | None = None, ) -> SimulationOptions: - """ - Set up the simulation options for the simulation API job. - """ - + """Set up the simulation options for the simulation API job.""" return SimulationOptions.model_validate( { "country": country_id, @@ -1647,13 +1620,11 @@ def _with_policyengine_bundle( return response def _setup_region(self, country_id: str, region: str) -> str: - """ - Validate the region for the given country. + """Validate the region for the given country. Assumes region has already been normalized (e.g., "ca" -> "state/ca"). Raises ValueError for invalid regions. """ - # For US regions, validate (skip validation for national "us") if country_id == "us" and region != "us": self._validate_us_region(region) @@ -1661,8 +1632,7 @@ def _setup_region(self, country_id: str, region: str) -> str: return region def _validate_us_region(self, region: str) -> None: - """ - Validate a prefixed US region string. + """Validate a prefixed US region string. Raises ValueError if the region is not valid. """ @@ -1740,9 +1710,7 @@ def _set_reform_impact_computing( setup_options: EconomicImpactSetupOptions, execution_id: str, ) -> None: - """ - In the reform_impact table, set the status of the impact to "computing". - """ + """In the reform_impact table, set the status of the impact to "computing".""" try: self._reform_impacts.set_reform_impact( country_id=setup_options.country_id, @@ -1778,9 +1746,7 @@ def _set_reform_impact_complete( reform_impact_json: dict[str, Any], execution_id: str, ) -> None: - """ - In the reform_impact table, set the status of the impact to "ok" and store the reform impact JSON. - """ + """In the reform_impact table, set the status of the impact to "ok" and store the reform impact JSON.""" try: self._reform_impacts.set_complete_reform_impact( country_id=setup_options.country_id, @@ -1812,9 +1778,7 @@ def _set_reform_impact_error( execution_id: str, error_code: str | None = None, ) -> None: - """ - In the reform_impact table, set the status of the impact to "error" and store the error message. - """ + """In the reform_impact table, set the status of the impact to "error" and store the error message.""" try: self._reform_impacts.set_error_reform_impact( country_id=setup_options.country_id, @@ -1839,5 +1803,4 @@ def _set_reform_impact_error( def _create_submission_claim_id(self) -> str: """Create an opaque token for one submission ownership claim.""" - return str(uuid.uuid4()) diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py index 655b7700b..67e99ff5a 100644 --- a/policyengine_api/services/household_calculation_service.py +++ b/policyengine_api/services/household_calculation_service.py @@ -362,7 +362,6 @@ def prepare_household_calculation( spm_requested: bool = False, ) -> PreparedHouseholdCalculation: """Validate request inputs before accepting a calculation.""" - countries = self._countries() country = countries.get(country_id) spm = normalize_spm_selection(country_id, spm) @@ -398,7 +397,6 @@ def calculate_prepared_household( prepared: PreparedHouseholdCalculation, ) -> HouseholdCalculationResult: """Calculate inputs already accepted by ``prepare_household_calculation``.""" - with observability_runtime.span( HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_CALCULATION) ): @@ -443,7 +441,6 @@ def parse_prepared_household( on_accepted: Callable[[], object] | None = None, ) -> PreparedHouseholdCalculation: """Parse a prepared situation without performing requested calculations.""" - prepare_calculation = getattr(prepared.country, "prepare_calculation", None) country_calculation = prepared.country_calculation if callable(prepare_calculation): @@ -476,7 +473,6 @@ def calculate_household( spm_requested: bool = False, ) -> HouseholdCalculationResult: """Validate and calculate request-provided household and policy data.""" - prepared = self.prepare_household_calculation( country_id, household_json, diff --git a/policyengine_api/services/household_mirroring.py b/policyengine_api/services/household_mirroring.py index 276e7ca90..940e9c4fa 100644 --- a/policyengine_api/services/household_mirroring.py +++ b/policyengine_api/services/household_mirroring.py @@ -124,7 +124,6 @@ def process_household_event_after_commit( require_pending: bool = False, ) -> LegacyHouseholdPersistenceResult: """Copy one exact retained create event and then record its completion.""" - selected_event_service = event_service or HouseholdService() started_at = time.perf_counter() destination_result: LegacyHouseholdPersistenceResult | None = None diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index 2957d5bd9..6fdd7ee68 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -278,7 +278,6 @@ def process_mirror_event( require_pending: bool = False, ) -> LegacyHouseholdPersistenceResult: """Process one explicitly selected retained create event.""" - with self._sessions.begin() as session: event = session.scalar( select(HouseholdMirrorEvent) diff --git a/policyengine_api/services/policy_mirroring.py b/policyengine_api/services/policy_mirroring.py index 257481f7f..217d2aba7 100644 --- a/policyengine_api/services/policy_mirroring.py +++ b/policyengine_api/services/policy_mirroring.py @@ -125,7 +125,6 @@ def mirror_policy_after_commit( mirror_factory: Callable[[], LegacyPolicyMirror] | None = None, ) -> LegacyPolicyPersistenceResult: """Require the Supabase policy and mapping transaction before v1 success.""" - started_at = time.perf_counter() try: result = (mirror_factory or _default_mirror_factory)().mirror_legacy_policy( diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index 5112d9429..714a59a38 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -31,7 +31,6 @@ class PolicySetResult: def __iter__(self) -> Iterator[int | str | bool]: """Preserve the established three-value internal unpacking interface.""" - yield self.policy_id yield self.message yield self.is_existing_policy @@ -95,7 +94,6 @@ def get_policy_snapshot( policy_id: int, ) -> LegacyPolicySnapshot | None: """Return detached fields required to mirror one existing v1 policy.""" - policy = self.get_policy(country_id, policy_id) if policy is None: return None diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 1658ff655..7a0c118f9 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -94,7 +94,6 @@ def claim_reform_impact_start( observability_id: str, ) -> bool: """Fail closed unless this request atomically owns job submission.""" - return self._cache.claim_start( country_id=country_id, reform_policy_id=policy_id, @@ -123,7 +122,6 @@ def get_reform_impact_start_claim( target: str, ) -> ReformImpactStartClaim | None: """Return the current submission owner and diagnostic identifier.""" - return self._cache.get_start_claim( country_id=country_id, reform_policy_id=policy_id, @@ -152,7 +150,6 @@ def release_reform_impact_start( observability_id: str, ) -> None: """Best-effort release; an unavailable cache safely falls back to expiry.""" - try: self._cache.release_start( country_id=country_id, diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index cbcfc7cf8..0749bac83 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -29,7 +29,6 @@ def _canonicalize_annual_spm_output( expected_year: str | None, ) -> JSONObject: """Validate and serialize compact SPM provenance on one economy result.""" - if "spm_config" in output: raise ValueError( "Completed US economy output must not include legacy spm_config" @@ -58,7 +57,6 @@ def _canonicalize_society_wide_spm_output( expected_year: str | None, ) -> object: """Validate compact annual or budget-window society-wide SPM output.""" - if not isinstance(output, dict): return output typed_output = cast("JSONObject", output) diff --git a/policyengine_api/services/user_policy_mirroring.py b/policyengine_api/services/user_policy_mirroring.py index 7f24d8de9..dddadf499 100644 --- a/policyengine_api/services/user_policy_mirroring.py +++ b/policyengine_api/services/user_policy_mirroring.py @@ -169,7 +169,6 @@ def mirror_user_policy_after_commit( mirror_factory: Callable[[], LegacyUserPolicyMirror] | None = None, ) -> LegacyUserPolicyPersistenceResult: """Require one complete Supabase association transaction before v1 success.""" - started_at = time.perf_counter() try: result = _run_user_policy_mirror( @@ -218,7 +217,6 @@ def mirror_pending_user_policy_events_after_commit( mirror_factory: Callable[[], LegacyUserPolicyMirror] | None = None, ) -> LegacyUserPolicyPersistenceResult: """Synchronously apply retained source events through one request's revision.""" - selected_event_service = event_service or UserPolicyService() selected_reform_loader = ( reform_snapshot_loader or PolicyService().get_policy_snapshot diff --git a/policyengine_api/services/user_policy_service.py b/policyengine_api/services/user_policy_service.py index 3818791e5..99ffc46ea 100644 --- a/policyengine_api/services/user_policy_service.py +++ b/policyengine_api/services/user_policy_service.py @@ -389,7 +389,6 @@ def process_pending_mirror_events( | None = None, ) -> LegacyUserPolicyPersistenceResult: """Process retained source events in order through one request's revision.""" - latest_result: LegacyUserPolicyPersistenceResult | None = None while True: with self._sessions.begin() as session: diff --git a/policyengine_api/services/v2/comparison_runs/validators.py b/policyengine_api/services/v2/comparison_runs/validators.py index 56dd42d8b..2e5fafa69 100644 --- a/policyengine_api/services/v2/comparison_runs/validators.py +++ b/policyengine_api/services/v2/comparison_runs/validators.py @@ -164,7 +164,6 @@ def require_simulation_parent_matches( child: ComparisonSimulationRecord, ) -> None: """Require a child to retain the parent's release and retention identity.""" - if child.evaluation_id != parent.evaluation_id: raise ComparisonRunIdentityError( "comparison child names a different parent evaluation_id" @@ -188,7 +187,6 @@ def require_successful_report_replay( candidate: ComparisonReportRecord, ) -> None: """Reject attempts to alter an already successful report receipt.""" - _require_equal_fields(existing, candidate, REPORT_SUCCEEDED_REPLAY_FIELDS) @@ -197,7 +195,6 @@ def require_successful_simulation_replay( candidate: ComparisonSimulationRecord, ) -> None: """Reject attempts to alter an already successful simulation receipt.""" - _require_equal_fields(existing, candidate, SIMULATION_SUCCEEDED_REPLAY_FIELDS) diff --git a/policyengine_api/services/v2/households/validators.py b/policyengine_api/services/v2/households/validators.py index 4f3e80fec..072faa299 100644 --- a/policyengine_api/services/v2/households/validators.py +++ b/policyengine_api/services/v2/households/validators.py @@ -52,7 +52,6 @@ def require_json_value( containers: frozenset[int] = frozenset(), ) -> object: """Reject values that are not standards-compliant JSON.""" - if depth > MAXIMUM_JSON_NESTING: raise ValueError("JSON values must not exceed 100 nested containers") if value is None or type(value) in {str, bool, int}: @@ -109,7 +108,6 @@ def normalize_household_document( document: dict[str, Any], ) -> dict[str, Any]: """Validate and normalize one country-specific household document.""" - collections = SUPPORTED_ENTITY_COLLECTIONS.get(country_id) if collections is None: raise HouseholdValidationError(f"unsupported country_id {country_id!r}") @@ -208,7 +206,6 @@ def verify_legacy_household_mapping( expected_household_id: UUID | None = None, ) -> None: """Validate one stored immutable legacy mapping without executing SQL.""" - if ( mapping.fingerprint_version != fingerprint_version or mapping.fingerprint_sha256 != fingerprint_sha256 diff --git a/policyengine_api/services/v2/metadata/validators.py b/policyengine_api/services/v2/metadata/validators.py index 6172d2018..45bb3c88f 100644 --- a/policyengine_api/services/v2/metadata/validators.py +++ b/policyengine_api/services/v2/metadata/validators.py @@ -48,7 +48,6 @@ def validate_economy_options( national_dataset: Dataset | None, ) -> Dataset: """Validate preloaded records required by the economy-options response.""" - if not any(region.code == country_id for region in regions): raise MetadataCatalogUnavailableError( f"the {country_id} national v2 region is absent" diff --git a/policyengine_api/services/v2/policies/validators.py b/policyengine_api/services/v2/policies/validators.py index 0ceabb046..d0ecf8120 100644 --- a/policyengine_api/services/v2/policies/validators.py +++ b/policyengine_api/services/v2/policies/validators.py @@ -52,7 +52,6 @@ def require_json_value( containers: frozenset[int] = frozenset(), ) -> object: """Reject values that cannot be represented as standards-compliant JSON.""" - if depth > MAXIMUM_JSON_NESTING: raise ValueError("JSON values must not exceed 100 nested containers") if value is None or type(value) in {str, bool, int}: @@ -80,7 +79,6 @@ def require_json_value( def normalize_utc(value: datetime) -> datetime: """Require a time-zone-aware value and normalize it to UTC.""" - if value.tzinfo is None or value.utcoffset() is None: raise ValueError("effective dates must include a UTC offset") return value.astimezone(UTC) @@ -93,7 +91,6 @@ def validate_policy_catalog( resolved_parameter_ids: set[UUID], ) -> ResolvedPolicyCreationInput: """Validate preloaded catalog records and bind them to policy content.""" - from policyengine_api.services.v2.policies.types import ResolvedPolicyCreationInput if policy_input.tax_benefit_model_id != selected.model.id: @@ -123,7 +120,6 @@ def verify_legacy_policy_mapping( expected_policy_id: UUID | None = None, ) -> None: """Validate an existing mapping without performing SQL.""" - if mapping.source_policy_hash != source_policy_hash: raise LegacyPolicyMappingIntegrityError( "legacy policy identity was presented with a different source hash" diff --git a/policyengine_api/services/v2/reports/types.py b/policyengine_api/services/v2/reports/types.py index ef4d997f4..8cb7041f1 100644 --- a/policyengine_api/services/v2/reports/types.py +++ b/policyengine_api/services/v2/reports/types.py @@ -107,7 +107,6 @@ def require_unique_aggregates( @model_validator(mode="after") def validate_spm_provenance(self) -> AggregateReportArtifactPayload: """Reject duplicate or noncanonical SPM data in aggregate artifacts.""" - if "spm_config" in self.result: raise ValueError("aggregate result must not include legacy spm_config") if "spm_provenance" in self.result: diff --git a/policyengine_api/services/v2/user_policies/transformations.py b/policyengine_api/services/v2/user_policies/transformations.py index 002987117..32f846e21 100644 --- a/policyengine_api/services/v2/user_policies/transformations.py +++ b/policyengine_api/services/v2/user_policies/transformations.py @@ -47,7 +47,6 @@ def user_policy_page( def fingerprint_legacy_user_policy(snapshot: LegacyUserPolicySnapshot) -> str: """Hash every committed source field through deterministic JSON.""" - document = { "fingerprint_version": USER_POLICY_FINGERPRINT_VERSION, **snapshot.model_dump(mode="json"), @@ -69,7 +68,6 @@ def project_legacy_user_policy( policy_id: UUID, ) -> UserPolicyCreationInput: """Map v1 presentation data onto an association, never core policy content.""" - return UserPolicyCreationInput( country_id=snapshot.country_id, user_id=user_id, diff --git a/policyengine_api/spm.py b/policyengine_api/spm.py index 9919aa729..ddfd78002 100644 --- a/policyengine_api/spm.py +++ b/policyengine_api/spm.py @@ -216,7 +216,6 @@ def build_spm_provenance( runtime_versions: SPMRuntimeVersions, ) -> SPMProvenance: """Build one compact receipt from typed calculation inputs.""" - if selection.forecast_content_sha256 != forecast_sha256: raise ValueError("SPM receipt differs from the SPM configuration") if selection.scenario is None: @@ -247,7 +246,6 @@ def build_spm_calculation_provenance( receipt: Mapping[str, object], ) -> SPMProvenance: """Reduce the country-owned rich receipt to the public compact contract.""" - forecast_sha256 = _required_string(receipt, "forecast_sha256") scenario = _required_string(receipt, "scenario") geography_kind = _required_string(receipt, "geography_kind") @@ -284,7 +282,6 @@ def resolved_spm_configuration( provenance: SPMProvenance, ) -> SPMResolvedConfiguration: """Derive the complete resolved selection from its compact receipt.""" - return SPMResolvedConfiguration( forecast_content_sha256=provenance.forecast_sha256, scenario=provenance.scenario, @@ -313,7 +310,6 @@ def build_spm_comparison_provenance( reform_receipts: Sequence[SPMProvenance], ) -> SPMComparisonProvenance: """Collapse matching child receipts into one comparison receipt.""" - return SPMComparisonProvenance( schema_version="canonical-spm-comparison-v2", baseline=_collapse_execution_receipts("baseline", baseline_receipts), @@ -425,7 +421,6 @@ def validate_spm_calculation_provenance( provenance: object, ) -> SPMProvenance: """Validate the sole transported compact receipt for a completed result.""" - return SPMProvenance.model_validate(provenance) diff --git a/policyengine_api/utils/cache_utils.py b/policyengine_api/utils/cache_utils.py index cbd02da9d..e977962f4 100644 --- a/policyengine_api/utils/cache_utils.py +++ b/policyengine_api/utils/cache_utils.py @@ -10,7 +10,7 @@ def make_cache_key(*args: object, **kwargs: object) -> str: # pylint: disable=unused-argument - """make a hash to uniquely identify a cache entry. + """Make a hash to uniquely identify a cache entry. keep it fast, adding overhead to try to add some minor chance of a cache hit is not worth it. diff --git a/policyengine_api/utils/input_validation.py b/policyengine_api/utils/input_validation.py index a74934197..3029c2760 100644 --- a/policyengine_api/utils/input_validation.py +++ b/policyengine_api/utils/input_validation.py @@ -51,7 +51,6 @@ def find_unrecognized_inputs( metadata: dict, ) -> list[UnrecognizedInput]: """Return calculate payload names that the country package cannot resolve.""" - return [ *find_unrecognized_household_inputs(household, metadata), *find_unrecognized_policy_inputs(policy, metadata), diff --git a/policyengine_api/utils/payload_validators/validate_country.py b/policyengine_api/utils/payload_validators/validate_country.py index 9fc4130c1..1ab713804 100644 --- a/policyengine_api/utils/payload_validators/validate_country.py +++ b/policyengine_api/utils/payload_validators/validate_country.py @@ -27,6 +27,7 @@ def validate_country( Returns: Response(400) if country is not valid, else continues + """ @wraps(func) diff --git a/policyengine_api/utils/payload_validators/validate_household_payload.py b/policyengine_api/utils/payload_validators/validate_household_payload.py index 19a80a672..73da22c84 100644 --- a/policyengine_api/utils/payload_validators/validate_household_payload.py +++ b/policyengine_api/utils/payload_validators/validate_household_payload.py @@ -9,14 +9,14 @@ def validate_household_payload( payload: Mapping[str, object], ) -> tuple[bool, str | None]: - """ - Validate the payload for a POST request to set a household's input data. + """Validate the payload for a POST request to set a household's input data. Args: payload (dict): The payload to validate. Returns: tuple[bool, str]: A tuple containing a boolean indicating whether the payload is valid and a message. + """ # Check that all required keys are present required_keys = ["data"] diff --git a/scripts/export_migration_contracts.py b/scripts/export_migration_contracts.py index 0c5e50711..04638a9eb 100644 --- a/scripts/export_migration_contracts.py +++ b/scripts/export_migration_contracts.py @@ -22,7 +22,6 @@ def build_payload() -> dict[str, Any]: """Build the migration contract payload from typed registries.""" - route_groups = [] for group in ROUTE_GROUPS: row = asdict(group) @@ -67,13 +66,11 @@ def build_payload() -> dict[str, Any]: def json_text(payload: dict[str, Any]) -> str: """Return the canonical JSON text for a migration contract payload.""" - return json.dumps(payload, indent=2, sort_keys=True) + "\n" def render_markdown(payload: dict[str, Any]) -> str: """Render migration contracts as reviewer- and agent-readable Markdown.""" - lines = [ "# Migration Contracts", "", @@ -142,7 +139,6 @@ def write_outputs( markdown_path: Path = DEFAULT_MARKDOWN, ) -> None: """Write generated migration contract artifacts.""" - json_path.parent.mkdir(parents=True, exist_ok=True) markdown_path.parent.mkdir(parents=True, exist_ok=True) json_path.write_text(json_text(payload)) diff --git a/scripts/process_v1_household_mirror_event.py b/scripts/process_v1_household_mirror_event.py index 8ffd8392e..c4e5d3c0b 100644 --- a/scripts/process_v1_household_mirror_event.py +++ b/scripts/process_v1_household_mirror_event.py @@ -45,7 +45,6 @@ def validate_command_environment( environ: Mapping[str, str], ) -> dict[str, str]: """Require matching source, destination, and operator declarations.""" - deployment_environment = _required(environ, "DEPLOYMENT_ENVIRONMENT") supabase_environment = _required(environ, "V2_SUPABASE_ENVIRONMENT") if deployment_environment != selected_environment: @@ -110,7 +109,6 @@ def _proxy_event_service( environ: Mapping[str, str], ) -> tuple[HouseholdService | None, Engine | None]: """Use the workflow-authenticated Cloud SQL proxy when it is selected.""" - raw_port = environ.get("POLICYENGINE_DB_PROXY_PORT", "").strip() if not raw_port: return None, None @@ -162,7 +160,6 @@ def _record_command_outcome( error_type: str | None = None, ) -> None: """Write an audit record without source data or credentials.""" - logger.log_struct( { "message": "Explicit v1 household create-event processing completed", diff --git a/scripts/rules/ci_scripts.py b/scripts/rules/ci_scripts.py index 5a0a6a066..071b2560f 100644 --- a/scripts/rules/ci_scripts.py +++ b/scripts/rules/ci_scripts.py @@ -44,7 +44,8 @@ class Findings: def multiline_runs(config: CiScripts, root: Path) -> set[str]: """``workflow :: job :: step`` for every run: block with more than one line of - code, or with commands chained by ``&&``, ``||`` or ``;``.""" + code, or with commands chained by ``&&``, ``||`` or ``;``. + """ found = set() for workflow in sorted(root.glob(config.workflows)): relative = workflow.relative_to(root).as_posix() diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py index e57194a3e..d2c71174f 100644 --- a/scripts/v1_database_migration.py +++ b/scripts/v1_database_migration.py @@ -43,7 +43,6 @@ def build_database_url( database: str, ) -> URL: """Build a SQLAlchemy URL without stringifying its credentials.""" - return URL.create( drivername="mysql+pymysql", username=username, @@ -71,7 +70,6 @@ def classify_database_state( def describe_metadata_difference(difference: object) -> str: """Return a stable, data-free description of an Alembic metadata diff.""" - item = difference if isinstance(item, list) and len(item) == 1: item = item[0] diff --git a/tests/conftest.py b/tests/conftest.py index 5bf715308..7b77c0c8f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,7 +21,6 @@ @pytest.fixture def mock_v2_runtime_settings(): """Avoid external v2 configuration access in ordinary API route tests.""" - with patch( "policyengine_api.data.v2.settings.load_v2_runtime_database_settings", return_value=object(), diff --git a/tests/env_variables/test_environment_variables.py b/tests/env_variables/test_environment_variables.py index 631f6baca..de4175d7b 100644 --- a/tests/env_variables/test_environment_variables.py +++ b/tests/env_variables/test_environment_variables.py @@ -27,7 +27,6 @@ class TestEnvironmentVariables: ) def test_hugging_face_token(self): """Test if HUGGING_FACE_TOKEN is valid by querying Hugging Face API.""" - token = os.getenv("HUGGING_FACE_TOKEN") assert token is not None, "HUGGING_FACE_TOKEN is not set" @@ -89,7 +88,6 @@ def test_hugging_face_token(self): ) def test_github_microdata_auth_token(self): """Test if POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN is valid by querying GitHub user API.""" - token = os.getenv("POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN") assert token is not None, "POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN is not set" diff --git a/tests/fixtures/integration/simulations.py b/tests/fixtures/integration/simulations.py index 22845e338..7ed5ace50 100644 --- a/tests/fixtures/integration/simulations.py +++ b/tests/fixtures/integration/simulations.py @@ -1,6 +1,4 @@ -""" -Test fixtures and constants for axes calculation tests -""" +"""Test fixtures and constants for axes calculation tests""" from policyengine_api.services.household_calculation_service import ( add_yearly_variables, diff --git a/tests/fixtures/libs/simulation_entrypoint.py b/tests/fixtures/libs/simulation_entrypoint.py index de5e508f2..c4451b052 100644 --- a/tests/fixtures/libs/simulation_entrypoint.py +++ b/tests/fixtures/libs/simulation_entrypoint.py @@ -1,5 +1,4 @@ -""" -Test fixtures for SimulationAPIModal. +"""Test fixtures for SimulationAPIModal. This module provides mock data, fixtures, and helper functions for testing the Modal simulation API client. @@ -163,8 +162,7 @@ def create_mock_httpx_response( json_data: dict = None, headers: dict | None = None, ): - """ - Helper function to create a mock httpx response. + """Helper function to create a mock httpx response. Parameters ---------- @@ -177,6 +175,7 @@ def create_mock_httpx_response( ------- MagicMock A mock httpx response object. + """ mock_response = MagicMock() mock_response.status_code = status_code @@ -211,8 +210,7 @@ def mock_modal_env_url(): @pytest.fixture def mock_httpx_client(): - """ - Mock httpx.Client for testing SimulationAPIModal. + """Mock httpx.Client for testing SimulationAPIModal. Returns a mock client that can be configured for different responses. """ diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index 4f3738d96..f4b99647b 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -265,8 +265,7 @@ def create_mock_modal_execution( error=None, policyengine_bundle=None, ): - """ - Helper function to create mock Modal execution objects. + """Helper function to create mock Modal execution objects. Parameters ---------- @@ -283,6 +282,7 @@ def create_mock_modal_execution( ------- MagicMock A mock ModalSimulationExecution object. + """ mock_execution = MagicMock() mock_execution.job_id = job_id diff --git a/tests/fixtures/v2_catalog.py b/tests/fixtures/v2_catalog.py index 948b9d0e7..6e69f789b 100644 --- a/tests/fixtures/v2_catalog.py +++ b/tests/fixtures/v2_catalog.py @@ -62,7 +62,6 @@ def _region( def bundle(*, policyengine_version: str = POLICYENGINE_VERSION) -> dict: """Return a minimal packaged-bundle manifest.""" - return { "bundle_version": policyengine_version, "policyengine_version": policyengine_version, @@ -158,7 +157,6 @@ def source_models( policyengine_version: str = POLICYENGINE_VERSION, ) -> dict[str, SimpleNamespace]: """Return US and UK public-model fixtures.""" - us_default_uri = "hf://policyengine/us/populace_us_2024.h5@fixture" uk_default_uri = "hf://policyengine/uk/enhanced_frs_2024_25.h5@fixture" return { @@ -218,7 +216,6 @@ def source_models( def installed_version(name: str) -> str: """Return fixture observed versions.""" - return DEPENDENCY_VERSIONS[name] @@ -227,7 +224,6 @@ def normalized_catalog( policyengine_version: str = POLICYENGINE_VERSION, ) -> NormalizedCatalog: """Return the complete deterministic normalized fixture.""" - return extract_catalog( bundle=bundle(policyengine_version=policyengine_version), policyengine_version=policyengine_version, diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index dc838e439..9f4b532f7 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -17,7 +17,6 @@ @pytest.fixture(scope="session") def disposable_v1_database_url() -> str: """Return a validated local MySQL URL for destructive integration tests.""" - from sqlalchemy.engine import make_url database_url = os.environ.get("ALEMBIC_DATABASE_URL", "") @@ -39,7 +38,6 @@ def disposable_v1_database_url() -> str: @pytest.fixture(scope="session") def disposable_v2_database_url() -> str: """Return a validated disposable PostgreSQL URL for integration tests.""" - from policyengine_api.data.v2.migration_target import ( V2_ALEMBIC_DISPOSABLE_TEST, load_v2_alembic_settings, diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index a4935ec1c..6817297b0 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -43,7 +43,6 @@ def _deployed_question_table() -> Table: """Describe the orphaned production table without adding it to ORM metadata.""" - metadata = MetaData() return Table( "question", diff --git a/tests/integration/test_live_v2_policies.py b/tests/integration/test_live_v2_policies.py index e1584c17a..368209ba1 100644 --- a/tests/integration/test_live_v2_policies.py +++ b/tests/integration/test_live_v2_policies.py @@ -34,7 +34,6 @@ def test_live_native_policy_and_user_policy_lifecycle( integration_probe_id: str, ) -> None: """Exercise every Phase 10 native route against the staging database.""" - country_id = "us" model = _one_catalog_item( api_client, diff --git a/tests/integration/test_runtime_cache_redis.py b/tests/integration/test_runtime_cache_redis.py index ba332c7b7..0f0144b06 100644 --- a/tests/integration/test_runtime_cache_redis.py +++ b/tests/integration/test_runtime_cache_redis.py @@ -96,7 +96,8 @@ def test_cached_household_calculation_is_shared_between_connections( def test_cached_calculation_receipts_survive_a_real_round_trip(redis_pair) -> None: """The selection is part of cache identity and its receipts travel with the result, so a replay under the same measurement reads back a complete receipt - and a replay under a different one misses.""" + and a replay under a different one misses. + """ first, second, namespace = redis_pair selection = { "forecast_content_sha256": "c" * 64, diff --git a/tests/integration/test_v1_schema_metadata_compatibility.py b/tests/integration/test_v1_schema_metadata_compatibility.py index ed16a0247..f08fdc6ae 100644 --- a/tests/integration/test_v1_schema_metadata_compatibility.py +++ b/tests/integration/test_v1_schema_metadata_compatibility.py @@ -10,7 +10,6 @@ def compare_existing_schema(database_url: str) -> list[str]: """Return metadata drift without mutating the target database.""" - engine = create_engine(database_url) try: with engine.connect() as connection: diff --git a/tests/to_refactor/conftest.py b/tests/to_refactor/conftest.py index fcd3b7580..b33f66ebf 100644 --- a/tests/to_refactor/conftest.py +++ b/tests/to_refactor/conftest.py @@ -32,7 +32,6 @@ def _seed_current_law_policies(engine: Engine) -> None: @pytest.fixture(scope="session", autouse=True) def isolated_legacy_v1_database() -> Iterator[None]: """Bind legacy tests to explicit in-memory persistence, never Cloud SQL.""" - engine = create_engine( "sqlite+pysqlite:///:memory:", connect_args={"check_same_thread": False}, diff --git a/tests/to_refactor/python/test_validate_country.py b/tests/to_refactor/python/test_validate_country.py index f3495ffe8..f51660729 100644 --- a/tests/to_refactor/python/test_validate_country.py +++ b/tests/to_refactor/python/test_validate_country.py @@ -5,16 +5,14 @@ @validate_country def foo(country_id, other): - """ - A simple dummy test method for validation testing. Must be defined outside of the class (or within + """A simple dummy test method for validation testing. Must be defined outside of the class (or within the test functions themselves) due to complications with the `self` parameter for class methods. """ return "bar" class TestValidateCountry: - """ - Test that the @validate_country decorator returns 404 if the country does not exist, otherwise + """Test that the @validate_country decorator returns 404 if the country does not exist, otherwise continues execution of the function. """ diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 7c8bd744e..7b2688dbf 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -26,7 +26,6 @@ def test_engine(): @pytest.fixture(autouse=True) def isolated_orm_database(test_engine, monkeypatch): """Bind runtime factories to a clean in-memory ORM database per test.""" - monkeypatch.setattr(orm, "get_v1_engine", lambda: test_engine) orm.clear_v1_session_factories() factory = orm.get_v1_session_factory() diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 5616b8dae..50d4cd1c5 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -1,5 +1,4 @@ -""" -Unit tests for SimulationAPIModal class. +"""Unit tests for SimulationAPIModal class. Tests the selectable simulation API HTTP client functionality including job submission, status polling, and error handling. diff --git a/tests/unit/routes/test_metadata_routes_streaming.py b/tests/unit/routes/test_metadata_routes_streaming.py index c9d8b5c9e..ed2423c14 100644 --- a/tests/unit/routes/test_metadata_routes_streaming.py +++ b/tests/unit/routes/test_metadata_routes_streaming.py @@ -28,7 +28,6 @@ def _load_metadata_blueprint_with_fake_service(): import, so the route module is loaded with the fake installed first and the module cache is restored afterwards. """ - sentinel = object() original_route_module = sys.modules.get(ROUTE_MODULE, sentinel) original_service_module = sys.modules.get(SERVICE_MODULE, sentinel) diff --git a/tests/unit/services/test_economy_comparison.py b/tests/unit/services/test_economy_comparison.py index 3982731b3..4e8c8e0a4 100644 --- a/tests/unit/services/test_economy_comparison.py +++ b/tests/unit/services/test_economy_comparison.py @@ -758,7 +758,8 @@ def test__income_change_formula_exact(self): class TestIntraDecileImpact: """Tests for the intra_decile_impact function — verifying correct - percentage change calculation and bucket assignment.""" + percentage change calculation and bucket assignment. + """ def test__5pct_gain_classified_below_5pct(self): """A uniform 5% income gain must NOT land in 'Gain more than 5%'. @@ -834,7 +835,8 @@ def test__no_change_classified_correctly(self): def test__near_zero_baseline_no_division_error(self): """Households with near-zero baseline income should not cause - division errors — the floor of 1 handles this.""" + division errors — the floor of 1 handles this. + """ baseline = _make_economy( incomes=[0.0] * 10, deciles=list(range(1, 11)), @@ -852,7 +854,8 @@ def test__near_zero_baseline_no_division_error(self): def test__zero_baseline_uses_floor_of_one(self): """When baseline income is 0, the max(B, 1) floor means the effective denominator is 1. A $0 -> $100 change should give - income_change = 100/1 = 100 (10000%), landing in >5%.""" + income_change = 100/1 = 100 (10000%), landing in >5%. + """ baseline = _make_economy( incomes=[0.0] * 10, deciles=list(range(1, 11)), @@ -873,7 +876,8 @@ def test__zero_baseline_uses_floor_of_one(self): def test__negative_baseline_handled(self): """Households with negative baseline income should be handled - by the max(B, 1) floor without producing NaN or Inf.""" + by the max(B, 1) floor without producing NaN or Inf. + """ baseline = _make_economy( incomes=[-500.0] * 10, deciles=list(range(1, 11)), @@ -890,7 +894,8 @@ def test__negative_baseline_handled(self): def test__2pct_gain_is_not_doubled(self): """A 2% gain stays in <5% even if doubled (2*2%=4% < 5%), so this - test alone would not catch the doubling bug.""" + test alone would not catch the doubling bug. + """ baseline = _make_economy( incomes=[50000.0] * 10, deciles=list(range(1, 11)), @@ -909,7 +914,8 @@ def test__2pct_gain_is_not_doubled(self): def test__4pct_gain_not_doubled_into_above_5pct(self): """A 4% gain must stay in <5%. With the doubling bug, 4% * 2 = 8% would incorrectly land in >5%. This is the tightest regression - test for the doubling bug on the gain side.""" + test for the doubling bug on the gain side. + """ baseline = _make_economy( incomes=[10000.0] * 10, deciles=list(range(1, 11)), @@ -944,7 +950,8 @@ def test__all_field_averages_deciles(self): class TestIntraWealthDecileImpact: """Tests for intra_wealth_decile_impact — same formula, keyed by - wealth decile instead of income decile.""" + wealth decile instead of income decile. + """ def test__5pct_gain_classified_below_5pct(self): """Regression test: 5% gain must not be doubled into >5% bucket.""" diff --git a/tests/unit/services/test_metadata_service.py b/tests/unit/services/test_metadata_service.py index a8e3bda7d..d52eff31b 100644 --- a/tests/unit/services/test_metadata_service.py +++ b/tests/unit/services/test_metadata_service.py @@ -65,14 +65,14 @@ def test_get_metadata_empty_country_id(self): def test_verify_metadata_for_given_country( self, country_id, current_law_id, test_regions ): - """ - Verifies metadata for a specific country contains expected values. + """Verifies metadata for a specific country contains expected values. Args: service: The MetadataService fixture country_id: Country identifier string current_law_id: Expected current law ID for this country test_regions: List of region codes that should be present for this country + """ # Get metadata for the country service = MetadataService() @@ -158,8 +158,7 @@ def test_default_dataset_uses_user_facing_label(self, country_id, expected_label ], ) def test_verify_region_types_for_given_country(self, country_id, expected_types): - """ - Verifies that all regions for UK and US have a 'type' field + """Verifies that all regions for UK and US have a 'type' field with valid values. """ service = MetadataService() diff --git a/tests/unit/test_country_spm.py b/tests/unit/test_country_spm.py index 3f618b7b0..5e4bab16f 100644 --- a/tests/unit/test_country_spm.py +++ b/tests/unit/test_country_spm.py @@ -95,7 +95,8 @@ def test_inherited_measurement_leaves_dependent_variables_unavailable( monkeypatch, code, axes ): """Nobody chose this measurement, so its dependants behave like any other - variable the model cannot compute: null, not a rejected request.""" + variable the model cannot compute: null, not a rejected request. + """ country, _ = _failing_country(monkeypatch, code) # The requested cell is seeded null, so returning at all is the observable # change; the sibling test below pins that the rest of the result survives. diff --git a/tests/unit/utils/test_cache_utils.py b/tests/unit/utils/test_cache_utils.py index ddf0fac1c..695b9d5d3 100644 --- a/tests/unit/utils/test_cache_utils.py +++ b/tests/unit/utils/test_cache_utils.py @@ -65,7 +65,8 @@ def test_make_cache_key_is_sha256_hex(): def test_make_cache_key_stable_across_processes(): """Two independent Python processes must produce the same cache key for the same inputs, even though they use different - PYTHONHASHSEED values.""" + PYTHONHASHSEED values. + """ script = textwrap.dedent( """ from flask import Flask From 0ca0460bca322acaf1d80c661897676cecdd0e17 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:49:25 +0400 Subject: [PATCH 02/12] Reword test docstrings for the pydocstyle rules Tests are exempt from the missing-docstring rules (D1) but not from the style rules, so give each flagged test and fixture docstring a one-line imperative summary followed by a blank line, describe the parameters the numpy-style sections left out, and correct the validate_country test class, which said 404 where the decorator returns 400. Co-Authored-By: Claude Opus 5.5 --- tests/fixtures/integration/simulations.py | 10 ++--- tests/fixtures/libs/simulation_entrypoint.py | 4 +- tests/fixtures/services/economy_service.py | 8 ++-- tests/integration/test_runtime_cache_redis.py | 4 +- tests/integration/test_simulations.py | 2 +- tests/to_refactor/python/test_error_routes.py | 12 +++--- .../python/test_validate_country.py | 13 ++++-- .../unit/services/test_economy_comparison.py | 43 +++++++++++-------- tests/unit/services/test_metadata_service.py | 4 +- tests/unit/test_country_spm.py | 6 ++- tests/unit/test_workflow_guard_scripts.py | 2 +- tests/unit/utils/test_cache_utils.py | 7 +-- 12 files changed, 68 insertions(+), 47 deletions(-) diff --git a/tests/fixtures/integration/simulations.py b/tests/fixtures/integration/simulations.py index 7ed5ace50..383976941 100644 --- a/tests/fixtures/integration/simulations.py +++ b/tests/fixtures/integration/simulations.py @@ -1,4 +1,4 @@ -"""Test fixtures and constants for axes calculation tests""" +"""Households, axes configurations and builders for the axes calculation tests.""" from policyengine_api.services.household_calculation_service import ( add_yearly_variables, @@ -43,21 +43,21 @@ # Pytest fixtures def create_base_household(): - """Returns a copy of the base household configuration""" + """Return a deep copy of BASE_HOUSEHOLD that a test may mutate freely.""" import copy return copy.deepcopy(BASE_HOUSEHOLD) def create_small_axes(): - """Returns the small axes configuration for faster tests""" + """Return a deep copy of the small axes configuration, which keeps tests fast.""" import copy return copy.deepcopy(SMALL_AXES_CONFIG) def create_household_with_axes(base_household, axes_config): - """Helper function to add axes to a household""" + """Return a deep copy of ``base_household`` with ``axes_config`` as its axes.""" import copy household = copy.deepcopy(base_household) @@ -66,6 +66,6 @@ def create_household_with_axes(base_household, axes_config): def setup_small_axes_household(base_household, small_axes_config): - """Fixture to setup a household with small axes for testing""" + """Add the axes and the test country's yearly variables to a base household.""" household_with_axes = create_household_with_axes(base_household, small_axes_config) return add_yearly_variables(household_with_axes, TEST_COUNTRY_ID) diff --git a/tests/fixtures/libs/simulation_entrypoint.py b/tests/fixtures/libs/simulation_entrypoint.py index c4451b052..6cbdc5961 100644 --- a/tests/fixtures/libs/simulation_entrypoint.py +++ b/tests/fixtures/libs/simulation_entrypoint.py @@ -162,7 +162,7 @@ def create_mock_httpx_response( json_data: dict = None, headers: dict | None = None, ): - """Helper function to create a mock httpx response. + """Build a mock httpx response for the Simulation Entrypoint client tests. Parameters ---------- @@ -170,6 +170,8 @@ def create_mock_httpx_response( HTTP status code for the response. json_data : dict JSON data to return from response.json(). + headers : dict or None + Response headers; defaults to a mock observability-id header. Returns ------- diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index f4b99647b..5330cc1ca 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -223,7 +223,7 @@ def create_mock_reform_impact( time_period=MOCK_TIME_PERIOD, message=None, ): - """Helper function to create mock reform impact records.""" + """Build a ReformImpact record populated with the shared mock identifiers.""" default_reform_impact_json = json.dumps( { **MOCK_REFORM_IMPACT_DATA, @@ -265,7 +265,7 @@ def create_mock_modal_execution( error=None, policyengine_bundle=None, ): - """Helper function to create mock Modal execution objects. + """Build a mock Modal simulation execution. Parameters ---------- @@ -277,6 +277,8 @@ def create_mock_modal_execution( The simulation result if complete. error : str or None The error message if failed. + policyengine_bundle : dict or None + The runtime bundle the execution reports; defaults to MOCK_RUNTIME_BUNDLE. Returns ------- @@ -306,7 +308,7 @@ def create_mock_budget_window_batch_execution( result=None, error=None, ): - """Helper function to create mock batch execution objects.""" + """Build a mock budget-window batch execution; year lists default to empty.""" mock_execution = MagicMock() mock_execution.batch_job_id = batch_job_id mock_execution.name = batch_job_id diff --git a/tests/integration/test_runtime_cache_redis.py b/tests/integration/test_runtime_cache_redis.py index 0f0144b06..941b4953d 100644 --- a/tests/integration/test_runtime_cache_redis.py +++ b/tests/integration/test_runtime_cache_redis.py @@ -94,7 +94,9 @@ def test_cached_household_calculation_is_shared_between_connections( def test_cached_calculation_receipts_survive_a_real_round_trip(redis_pair) -> None: - """The selection is part of cache identity and its receipts travel with the + """Key cached calculations by SPM selection and keep their receipts. + + The selection is part of cache identity and its receipts travel with the result, so a replay under the same measurement reads back a complete receipt and a replay under a different one misses. """ diff --git a/tests/integration/test_simulations.py b/tests/integration/test_simulations.py index 7bc19e2e3..64c53aa2e 100644 --- a/tests/integration/test_simulations.py +++ b/tests/integration/test_simulations.py @@ -14,7 +14,7 @@ class TestSimsWithAxes: def test__given_any_number_of_axes__sim_returns_valid_arrays( self, ): # , patched_add_yearly_variables, patched_countries_get): - """Integration test with small axes for speed""" + """Calculate a household with small axes, which keeps the test fast.""" base_household = create_base_household() small_axes_config = create_small_axes() household_with_axes = setup_small_axes_household( diff --git a/tests/to_refactor/python/test_error_routes.py b/tests/to_refactor/python/test_error_routes.py index 09c2058da..34b24fbc4 100644 --- a/tests/to_refactor/python/test_error_routes.py +++ b/tests/to_refactor/python/test_error_routes.py @@ -26,7 +26,7 @@ def client(app): def test_404_handler(app, client): - """Test 404 Not Found error handling""" + """Map NotFound to a 404 error envelope that keeps its message.""" @app.route("/nonexistent") def nonexistent(): @@ -42,7 +42,7 @@ def nonexistent(): def test_400_handler(app, client): - """Test 400 Bad Request error handling""" + """Map BadRequest to a 400 error envelope that keeps its message.""" @app.route("/bad-request") def bad_request(): @@ -58,7 +58,7 @@ def bad_request(): def test_401_handler(app, client): - """Test 401 Unauthorized error handling""" + """Map Unauthorized to a 401 error envelope that keeps its message.""" @app.route("/unauthorized") def unauthorized(): @@ -74,7 +74,7 @@ def unauthorized(): def test_403_handler(app, client): - """Test 403 Forbidden error handling""" + """Map Forbidden to a 403 error envelope that keeps its message.""" @app.route("/forbidden") def forbidden(): @@ -90,7 +90,7 @@ def forbidden(): def test_500_handler(app, client): - """Test 500 Internal Server Error handling""" + """Map InternalServerError to a 500 error envelope that keeps its message.""" @app.route("/server-error") def server_error(): @@ -106,7 +106,7 @@ def server_error(): def test_generic_exception_handler(app, client): - """Test handling of generic exceptions""" + """Map any other exception to a 500 error envelope that keeps its message.""" @app.route("/generic-error") def generic_error(): diff --git a/tests/to_refactor/python/test_validate_country.py b/tests/to_refactor/python/test_validate_country.py index f51660729..389603d37 100644 --- a/tests/to_refactor/python/test_validate_country.py +++ b/tests/to_refactor/python/test_validate_country.py @@ -5,15 +5,20 @@ @validate_country def foo(country_id, other): - """A simple dummy test method for validation testing. Must be defined outside of the class (or within - the test functions themselves) due to complications with the `self` parameter for class methods. + """Return "bar"; a stand-in view for exercising the decorator. + + It must be defined outside the class (or within the test functions themselves) + because the decorator treats the first argument as the country ID, and on a + method that would be `self`. """ return "bar" class TestValidateCountry: - """Test that the @validate_country decorator returns 404 if the country does not exist, otherwise - continues execution of the function. + """Tests for the @validate_country decorator. + + The decorator returns a 400 response if the country does not exist, and + otherwise calls the wrapped function. """ def test_valid_country(self): diff --git a/tests/unit/services/test_economy_comparison.py b/tests/unit/services/test_economy_comparison.py index 4e8c8e0a4..07eba32dc 100644 --- a/tests/unit/services/test_economy_comparison.py +++ b/tests/unit/services/test_economy_comparison.py @@ -738,7 +738,7 @@ def _make_economy( people=None, decile_key="household_income_decile", ): - """Helper to build baseline/reform dicts for intra_decile tests.""" + """Build a baseline or reform economy dict for the intra-decile tests.""" n = len(incomes) return { "household_net_income": np.array(incomes, dtype=float), @@ -757,8 +757,9 @@ def test__income_change_formula_exact(self): class TestIntraDecileImpact: - """Tests for the intra_decile_impact function — verifying correct - percentage change calculation and bucket assignment. + """Tests for intra_decile_impact. + + They verify the percentage change calculation and the bucket assignment. """ def test__5pct_gain_classified_below_5pct(self): @@ -834,8 +835,9 @@ def test__no_change_classified_correctly(self): assert pct == 1.0 def test__near_zero_baseline_no_division_error(self): - """Households with near-zero baseline income should not cause - division errors — the floor of 1 handles this. + """Avoid division errors for households with zero baseline income. + + The floor of 1 on the denominator handles this. """ baseline = _make_economy( incomes=[0.0] * 10, @@ -852,9 +854,10 @@ def test__near_zero_baseline_no_division_error(self): assert abs(total - 1.0) < 1e-9, f"Proportions should sum to 1, got {total}" def test__zero_baseline_uses_floor_of_one(self): - """When baseline income is 0, the max(B, 1) floor means the - effective denominator is 1. A $0 -> $100 change should give - income_change = 100/1 = 100 (10000%), landing in >5%. + """Divide by the floor of 1 when baseline income is 0. + + The max(B, 1) floor makes the effective denominator 1, so a $0 -> $100 + change gives income_change = 100/1 = 100 (10000%), landing in >5%. """ baseline = _make_economy( incomes=[0.0] * 10, @@ -875,8 +878,9 @@ def test__zero_baseline_uses_floor_of_one(self): assert not np.isinf(result["all"][label]) def test__negative_baseline_handled(self): - """Households with negative baseline income should be handled - by the max(B, 1) floor without producing NaN or Inf. + """Handle negative baseline income without producing NaN or Inf. + + The max(B, 1) floor applies to negative incomes too. """ baseline = _make_economy( incomes=[-500.0] * 10, @@ -893,8 +897,10 @@ def test__negative_baseline_handled(self): assert not np.isinf(result["all"][label]) def test__2pct_gain_is_not_doubled(self): - """A 2% gain stays in <5% even if doubled (2*2%=4% < 5%), so this - test alone would not catch the doubling bug. + """Keep a 2% gain in the <5% bucket. + + A 2% gain stays in <5% even if doubled (2*2%=4% < 5%), so this test + alone would not catch the doubling bug. """ baseline = _make_economy( incomes=[50000.0] * 10, @@ -912,9 +918,10 @@ def test__2pct_gain_is_not_doubled(self): assert pct == 1.0, "2% gain not classified as <5%" def test__4pct_gain_not_doubled_into_above_5pct(self): - """A 4% gain must stay in <5%. With the doubling bug, 4% * 2 = 8% - would incorrectly land in >5%. This is the tightest regression - test for the doubling bug on the gain side. + """Keep a 4% gain in the <5% bucket. + + With the doubling bug, 4% * 2 = 8% would incorrectly land in >5%. This + is the tightest regression test for the doubling bug on the gain side. """ baseline = _make_economy( incomes=[10000.0] * 10, @@ -949,8 +956,10 @@ def test__all_field_averages_deciles(self): class TestIntraWealthDecileImpact: - """Tests for intra_wealth_decile_impact — same formula, keyed by - wealth decile instead of income decile. + """Tests for intra_wealth_decile_impact. + + It uses the same formula as intra_decile_impact, keyed by wealth decile + instead of income decile. """ def test__5pct_gain_classified_below_5pct(self): diff --git a/tests/unit/services/test_metadata_service.py b/tests/unit/services/test_metadata_service.py index d52eff31b..25927e75f 100644 --- a/tests/unit/services/test_metadata_service.py +++ b/tests/unit/services/test_metadata_service.py @@ -158,9 +158,7 @@ def test_default_dataset_uses_user_facing_label(self, country_id, expected_label ], ) def test_verify_region_types_for_given_country(self, country_id, expected_types): - """Verifies that all regions for UK and US have a 'type' field - with valid values. - """ + """Give every UK and US region a 'type' field with a valid value.""" service = MetadataService() metadata = service.get_metadata(country_id) diff --git a/tests/unit/test_country_spm.py b/tests/unit/test_country_spm.py index 5e4bab16f..2b7666539 100644 --- a/tests/unit/test_country_spm.py +++ b/tests/unit/test_country_spm.py @@ -46,7 +46,7 @@ def requested_household(variable="spm_unit_spm_threshold", *, axes=False): def _country_raising(monkeypatch, error): - """A country whose every variable calculation raises `error`.""" + """Return a country whose every variable calculation raises `error`.""" def calculate(variable, period): # This also represents a requested resource with a nested SPM dependency. @@ -94,7 +94,9 @@ def test_country_never_swallows_spm_input_errors_for_a_chosen_measurement( def test_inherited_measurement_leaves_dependent_variables_unavailable( monkeypatch, code, axes ): - """Nobody chose this measurement, so its dependants behave like any other + """Leave variables that depend on an inherited SPM measurement null. + + Nobody chose this measurement, so its dependants behave like any other variable the model cannot compute: null, not a rejected request. """ country, _ = _failing_country(monkeypatch, code) diff --git a/tests/unit/test_workflow_guard_scripts.py b/tests/unit/test_workflow_guard_scripts.py index 9283cc15b..0763e7f15 100644 --- a/tests/unit/test_workflow_guard_scripts.py +++ b/tests/unit/test_workflow_guard_scripts.py @@ -49,7 +49,7 @@ def test_require_git_ref_fails_with_an_annotation_on_another_ref() -> None: def _fake_gcloud(tmp_path: Path, rows: list[str], *, delete_exit: int = 0) -> Path: - """A gcloud that lists ``rows`` and records every other call.""" + """Return a bin dir whose fake gcloud lists ``rows`` and logs every other call.""" bin_dir = tmp_path / "bin" bin_dir.mkdir() listing = tmp_path / "listing.txt" diff --git a/tests/unit/utils/test_cache_utils.py b/tests/unit/utils/test_cache_utils.py index 695b9d5d3..0b75db3ad 100644 --- a/tests/unit/utils/test_cache_utils.py +++ b/tests/unit/utils/test_cache_utils.py @@ -63,9 +63,10 @@ def test_make_cache_key_is_sha256_hex(): def test_make_cache_key_stable_across_processes(): - """Two independent Python processes must produce the same cache - key for the same inputs, even though they use different - PYTHONHASHSEED values. + """Produce the same cache key in two independent Python processes. + + The processes use different PYTHONHASHSEED values, so the key must not + depend on Python's randomised string hashing. """ script = textwrap.dedent( """ From 3fae045bb77d35b2d8f34ad15e83f324fa4a4cc4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:52:35 +0400 Subject: [PATCH 03/12] Document the runtime cache, observability and gateway libraries Add docstrings to the Redis-backed runtime cache (its Protocols, the in-memory fake, claims and the household and reform-impact caches), the observability runtime and stage plans, and the gateway client and auth helpers. Each says what a reader cannot see from the signature: what a miss means, when entries are rejected, and which Redis behaviour a fake command emulates. Adding a docstring to the fake backend's get makes RET501 report its explicit return None, which stays and is suppressed with a reason. Co-Authored-By: Claude Opus 5.5 --- policyengine_api/libs/gateway_auth.py | 8 +- .../libs/simulation_entrypoint.py | 12 +++ policyengine_api/observability/runtime.py | 1 + policyengine_api/observability/stages.py | 7 ++ policyengine_api/runtime_cache/claims.py | 13 +++ policyengine_api/runtime_cache/client.py | 5 ++ policyengine_api/runtime_cache/core.py | 88 ++++++++++++++++--- .../runtime_cache/dependencies.py | 7 ++ policyengine_api/runtime_cache/fake.py | 53 ++++++++++- .../runtime_cache/household_calculations.py | 21 +++++ .../runtime_cache/reform_impacts.py | 30 +++++++ policyengine_api/runtime_cache/settings.py | 3 + 12 files changed, 232 insertions(+), 16 deletions(-) diff --git a/policyengine_api/libs/gateway_auth.py b/policyengine_api/libs/gateway_auth.py index 8c4b28252..dcdc0dde7 100644 --- a/policyengine_api/libs/gateway_auth.py +++ b/policyengine_api/libs/gateway_auth.py @@ -49,7 +49,10 @@ class GatewayAuthError(RuntimeError): def gateway_auth_required() -> bool: - """True iff this runtime requires gateway auth to be configured.""" + """Return whether ``GATEWAY_AUTH_REQUIRED`` makes missing gateway auth an error. + + Only ``1``, ``true``, ``yes`` and ``on`` (any case) enable it. + """ return os.environ.get(GATEWAY_AUTH_REQUIRED_ENV, "").lower() in ( "1", "true", @@ -120,6 +123,7 @@ def __init__( *, http_timeout: float = 10.0, ) -> None: + """Read each Auth0 setting from its argument, else its environment variable.""" self._issuer = ( issuer if issuer is not None @@ -256,11 +260,13 @@ class GatewayBearerAuth(httpx.Auth): """Attach a bearer token and retry once on a 401.""" def __init__(self, token_provider: GatewayAuthTokenProvider) -> None: + """Authenticate requests with tokens from ``token_provider``.""" self._token_provider = token_provider def auth_flow( self, request: httpx.Request ) -> Generator[httpx.Request, httpx.Response, None]: + """Add a bearer token; on a 401, refetch the token and resend once.""" request.headers["Authorization"] = f"Bearer {self._token_provider.get_token()}" response = yield request diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 1eead55d5..ecea6e8e6 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -113,6 +113,11 @@ class SimulationEntrypointClient: """ def __init__(self, entrypoint: str | None = None) -> None: + """Connect to ``entrypoint`` (default: the configured one) with gateway auth. + + Raise GatewayAuthError if auth is partly configured, or is required here + but not configured. + """ self.entrypoint = entrypoint or get_sim_entrypoint() self.base_url = resolve_simulation_entrypoint_url(self.entrypoint) self._token_provider = GatewayAuthTokenProvider() @@ -307,6 +312,13 @@ def get_spm_capability( *, policyengine_version: str | None = None, ) -> object: + """Return the worker's published SPM capability for the bundle it would run. + + The bundle is the one the worker's ``/versions`` routes resolve for + ``policyengine_version`` or else ``country``/``version``; None means that + bundle publishes no capability. Raise SPMValidationError if the routes + are unreachable or ambiguous. + """ from policyengine_api.spm import SPMValidationError try: diff --git a/policyengine_api/observability/runtime.py b/policyengine_api/observability/runtime.py index b944ab3eb..a560b5378 100644 --- a/policyengine_api/observability/runtime.py +++ b/policyengine_api/observability/runtime.py @@ -68,6 +68,7 @@ def _build_runtime() -> ObservabilityRuntime: def get_runtime() -> ObservabilityRuntime: + """Return the process-wide runtime that was configured when this module loaded.""" return runtime diff --git a/policyengine_api/observability/stages.py b/policyengine_api/observability/stages.py index 2c517542b..698a69475 100644 --- a/policyengine_api/observability/stages.py +++ b/policyengine_api/observability/stages.py @@ -12,12 +12,16 @@ class RunConfiguration(StrEnum): + """A kind of calculation request, each with its own registered stage plan.""" + HOUSEHOLD = "household" ECONOMY_ANNUAL = "economy_annual" ECONOMY_BUDGET_WINDOW = "economy_budget_window" class Stage(StrEnum): + """A named calculation step; each value is the span name telemetry records.""" + HOUSEHOLD_LOAD_INPUTS = "household.load_inputs" HOUSEHOLD_CACHE_LOOKUP = "household.cache_lookup" HOUSEHOLD_INPUT_NORMALIZATION = "household.input_normalization" @@ -43,10 +47,13 @@ class Stage(StrEnum): @dataclass(frozen=True) class StagePlan: + """The stages one run configuration is allowed to record as spans.""" + configuration: RunConfiguration stages: tuple[Stage, ...] def name(self, stage: Stage) -> str: + """Return the span name for ``stage``; raise ValueError if not in this plan.""" if stage not in self.stages: raise ValueError( f"{stage.value!r} is not registered for {self.configuration.value!r}" diff --git a/policyengine_api/runtime_cache/claims.py b/policyengine_api/runtime_cache/claims.py index 3fcdc58ab..7ce7008f0 100644 --- a/policyengine_api/runtime_cache/claims.py +++ b/policyengine_api/runtime_cache/claims.py @@ -19,16 +19,24 @@ class ExpiringClaimStore: + """Token-owned keys that let one worker claim shared work until a TTL lapses.""" + def __init__( self, client: CacheBackend, *, family: str = "coordination", ) -> None: + """Store claims in ``client``, reporting events under the ``family`` label.""" self.client = client self.family = family def acquire(self, key: str, token: str, *, ttl_seconds: int) -> bool: + """Claim ``key`` for ``token`` if it is free, returning whether it succeeded. + + Raise CacheCoordinationError when the backend fails, so callers never + proceed as if they owned the work. + """ if not token or ttl_seconds <= 0: raise ValueError("claim token and positive TTL are required") started_at = time.perf_counter() @@ -54,6 +62,11 @@ def acquire(self, key: str, token: str, *, ttl_seconds: int) -> bool: return acquired def release(self, key: str, token: str) -> bool: + """Delete the claim only if ``token`` still owns it; return whether it did. + + A claim that expired or passed to another owner is left alone. Raise + CacheCoordinationError when the backend fails. + """ started_at = time.perf_counter() try: released = bool(self.client.eval(COMPARE_AND_DELETE, 1, key, token)) diff --git a/policyengine_api/runtime_cache/client.py b/policyengine_api/runtime_cache/client.py index a7225db18..13ad1cc51 100644 --- a/policyengine_api/runtime_cache/client.py +++ b/policyengine_api/runtime_cache/client.py @@ -18,6 +18,10 @@ def build_runtime_cache_client(settings: RuntimeCacheSettings) -> redis.Redis: + """Build a Redis client with bounded timeouts and no retry on timeout. + + Raise RuntimeError when the cache is disabled or TLS has no CA bundle. + """ if not settings.enabled or settings.url is None: raise RuntimeError("the shared runtime cache is disabled") kwargs: dict[str, Any] = { @@ -51,6 +55,7 @@ def get_runtime_cache_client( def close_runtime_cache_clients() -> None: + """Close and forget every cached client; the ASGI app calls this at shutdown.""" for client in _clients.values(): client.close() _clients.clear() diff --git a/policyengine_api/runtime_cache/core.py b/policyengine_api/runtime_cache/core.py index 058c31b9d..f6bc712ea 100644 --- a/policyengine_api/runtime_cache/core.py +++ b/policyengine_api/runtime_cache/core.py @@ -21,23 +21,65 @@ class CachePipeline(Protocol): """The queued commands that cache writers issue inside one pipeline.""" - def __enter__(self) -> Self: ... - def __exit__(self, *args: object) -> object: ... - def set(self, key: str, value: str, **kwargs: object) -> object: ... - def delete(self, *keys: str) -> object: ... - def expire(self, key: str, seconds: int) -> object: ... - def zadd(self, key: str, mapping: dict[str, float]) -> object: ... - def zrem(self, key: str, *members: str) -> object: ... - def zremrangebyrank(self, key: str, start: int, end: int) -> object: ... - def execute(self) -> list[object]: ... + def __enter__(self) -> Self: + """Open the pipeline for queuing commands.""" + ... + + def __exit__(self, *args: object) -> object: + """Discard any commands that were queued but not executed.""" + ... + + def set(self, key: str, value: str, **kwargs: object) -> object: + """Queue a Redis SET; ``kwargs`` carries options such as ``ex``.""" + ... + + def delete(self, *keys: str) -> object: + """Queue a Redis DEL of every given key.""" + ... + + def expire(self, key: str, seconds: int) -> object: + """Queue a Redis EXPIRE that sets the key's TTL in seconds.""" + ... + + def zadd(self, key: str, mapping: dict[str, float]) -> object: + """Queue a Redis ZADD of member-to-score pairs into a sorted set.""" + ... + + def zrem(self, key: str, *members: str) -> object: + """Queue a Redis ZREM of the given sorted-set members.""" + ... + + def zremrangebyrank(self, key: str, start: int, end: int) -> object: + """Queue a Redis ZREMRANGEBYRANK over inclusive ranks (``-1`` is last).""" + ... + + def execute(self) -> list[object]: + """Run the queued commands and return their replies in queue order.""" + ... class CacheBackend(Protocol): - def get(self, key: str) -> str | bytes | None: ... - def set(self, key: str, value: str, **kwargs: object) -> object: ... - def delete(self, *keys: str) -> object: ... - def pipeline(self, transaction: bool = True) -> CachePipeline: ... - def eval(self, script: str, numkeys: int, *keys_and_args: str) -> object: ... + """The Redis client subset the runtime cache uses; fakes implement it too.""" + + def get(self, key: str) -> str | bytes | None: + """Return the stored value, or None when the key is missing or expired.""" + ... + + def set(self, key: str, value: str, **kwargs: object) -> object: + """Store a value; ``kwargs`` carries Redis options such as ``ex`` and ``nx``.""" + ... + + def delete(self, *keys: str) -> object: + """Delete every given key and return how many existed.""" + ... + + def pipeline(self, transaction: bool = True) -> CachePipeline: + """Return a pipeline whose commands run together on ``execute``.""" + ... + + def eval(self, script: str, numkeys: int, *keys_and_args: str) -> object: + """Run a Lua script atomically on the server against the given keys.""" + ... class CacheCoordinationError(RuntimeError): @@ -46,6 +88,8 @@ class CacheCoordinationError(RuntimeError): @dataclass(frozen=True) class CacheNamespace: + """The environment and service prefix that keeps deployments' keys apart.""" + environment: str service: str @@ -57,6 +101,11 @@ def key( *, suffix: str | None = None, ) -> str: + """Build a key from a SHA-256 of the canonical JSON of ``inputs``. + + Equal inputs give the same key whatever their dict order; values that + are not JSON-serialisable are hashed by their ``str()``. + """ encoded = json.dumps( inputs, sort_keys=True, @@ -71,6 +120,7 @@ def key( return f"{key}:{suffix}" if suffix else key def family_key(self, family: str, schema_version: int, name: str) -> str: + """Build an unhashed key for a fixed, named entry within a cache family.""" return ( f"{CACHE_KEY_ROOT}:{self.environment}:{self.service}:" f"{family}:v{schema_version}:{name}" @@ -82,6 +132,7 @@ def encode_envelope( schema_version: int, payload: object, ) -> str: + """Serialise a payload with the family and schema version it was written under.""" return json.dumps( { "family": family, @@ -100,6 +151,11 @@ def decode_envelope( family: str, schema_version: int, ) -> object: + """Return an envelope's payload, or None if it is absent or unusable. + + Undecodable JSON and envelopes from another family or schema version are + treated as misses rather than errors. + """ if value is None: return None try: @@ -172,6 +228,7 @@ def __init__( schema_version: int, ttl_seconds: int, ) -> None: + """Bind the cache to one key family; raise ValueError for a non-positive TTL.""" if ttl_seconds <= 0: raise ValueError("cache TTL must be positive") self.client = client @@ -181,9 +238,11 @@ def __init__( self.ttl_seconds = ttl_seconds def key(self, inputs: dict[str, Any]) -> str: + """Return the namespaced key this cache stores ``inputs`` under.""" return self.namespace.key(self.family, self.schema_version, inputs) def get(self, inputs: dict[str, Any]) -> object: + """Return the cached payload, or None on a miss or any backend failure.""" started_at = time.perf_counter() try: value = self.client.get(self.key(inputs)) @@ -208,6 +267,7 @@ def get(self, inputs: dict[str, Any]) -> object: return payload def set(self, inputs: dict[str, Any], payload: object) -> bool: + """Store a payload with a jittered TTL; return False if the write failed.""" started_at = time.perf_counter() try: result = self.client.set( diff --git a/policyengine_api/runtime_cache/dependencies.py b/policyengine_api/runtime_cache/dependencies.py index c5cf2a97d..17324ef22 100644 --- a/policyengine_api/runtime_cache/dependencies.py +++ b/policyengine_api/runtime_cache/dependencies.py @@ -13,6 +13,8 @@ @dataclass(frozen=True) class RuntimeCacheContext: + """The process's cache backend and key namespace, built once from settings.""" + client: CacheBackend namespace: CacheNamespace enabled: bool @@ -20,6 +22,10 @@ class RuntimeCacheContext: @lru_cache(maxsize=1) def get_runtime_cache_context() -> RuntimeCacheContext: + """Return the shared cache context, a no-storage backend when disabled. + + Built on first use, not at import, and memoised for the process. + """ settings = load_runtime_cache_settings() namespace = CacheNamespace(settings.environment, settings.service) client: CacheBackend = ( @@ -35,4 +41,5 @@ def get_runtime_cache_context() -> RuntimeCacheContext: def clear_runtime_cache_context() -> None: + """Forget the memoised context so the next call rereads the settings.""" get_runtime_cache_context.cache_clear() diff --git a/policyengine_api/runtime_cache/fake.py b/policyengine_api/runtime_cache/fake.py index 74fe06318..f7c50f834 100644 --- a/policyengine_api/runtime_cache/fake.py +++ b/policyengine_api/runtime_cache/fake.py @@ -13,6 +13,7 @@ class InMemoryCacheBackend: """Small Redis-like fake with explicit time advancement and atomic pipelines.""" def __init__(self, *, now: Callable[[], float] | None = None) -> None: + """Start empty at time zero, or follow the ``now`` clock when one is given.""" self._external_now = now self._time = 0.0 self._values: dict[str, str] = {} @@ -24,6 +25,10 @@ def _now(self) -> float: return self._external_now() if self._external_now else self._time def advance(self, seconds: float) -> None: + """Move the fake clock forward so TTLs can expire without sleeping. + + Raises RuntimeError when the fake follows an external clock instead. + """ if self._external_now is not None: raise RuntimeError("cannot advance an externally clocked fake") self._time += seconds @@ -36,6 +41,7 @@ def _purge(self, key: str) -> None: self._expires.pop(key, None) def get(self, key: str) -> str | None: + """Emulate GET: return the string value, or None if missing or expired.""" with self._lock: self._purge(key) return self._values.get(key) @@ -48,6 +54,11 @@ def set( ex: int | None = None, nx: bool = False, ) -> bool | None: + """Emulate SET with the ``ex`` and ``nx`` options only. + + Like redis-py, return True on a write and None when ``nx`` finds the key + already set; a write without ``ex`` clears any earlier TTL. + """ with self._lock: self._purge(key) if nx and key in self._values: @@ -60,6 +71,7 @@ def set( return True def delete(self, *keys: str) -> int: + """Emulate DEL for strings and sorted sets, returning how many existed.""" with self._lock: deleted = 0 for key in keys: @@ -72,6 +84,7 @@ def delete(self, *keys: str) -> int: return deleted def expire(self, key: str, seconds: int) -> bool: + """Emulate EXPIRE: set a TTL and return False if the key does not exist.""" with self._lock: self._purge(key) if key not in self._values and key not in self._sorted_sets: @@ -80,6 +93,7 @@ def expire(self, key: str, seconds: int) -> bool: return True def zadd(self, key: str, mapping: dict[str, float]) -> int: + """Emulate ZADD without flags: upsert scores and count new members only.""" with self._lock: self._purge(key) values = self._sorted_sets.setdefault(key, {}) @@ -88,6 +102,11 @@ def zadd(self, key: str, mapping: dict[str, float]) -> int: return added def zrevrange(self, key: str, start: int, end: int) -> list[str]: + """Emulate ZREVRANGE: members by descending score over inclusive ranks. + + Equal scores fall back to reverse member order, as in Redis. The only + negative ``end`` supported is ``-1`` (the last member). + """ with self._lock: self._purge(key) ordered = sorted( @@ -99,6 +118,7 @@ def zrevrange(self, key: str, start: int, end: int) -> list[str]: return [member for member, _ in ordered[start:stop]] def zrem(self, key: str, *members: str) -> int: + """Emulate ZREM, returning how many of the members were present.""" with self._lock: values = self._sorted_sets.get(key, {}) removed = sum(member in values for member in members) @@ -107,6 +127,10 @@ def zrem(self, key: str, *members: str) -> int: return removed def zremrangebyrank(self, key: str, start: int, end: int) -> int: + """Emulate ZREMRANGEBYRANK over ascending-score ranks, inclusive of ``end``. + + As with ``zrevrange``, the only negative ``end`` supported is ``-1``. + """ with self._lock: self._purge(key) ordered = sorted( @@ -123,6 +147,11 @@ def eval( numkeys: int, *keys_and_args: str, ) -> int: + """Emulate EVAL of the claim store's compare-and-delete script only. + + The script text is ignored: delete the one key if it still holds the + expected token, and raise ValueError for any other key/argument shape. + """ if numkeys != 1 or len(keys_and_args) != 2: raise ValueError("fake supports one-key compare-and-delete only") key, expected = keys_and_args @@ -132,16 +161,22 @@ def eval( return self.delete(key) def pipeline(self, transaction: bool = True) -> InMemoryPipeline: + """Return a pipeline that queues commands until ``execute``.""" return InMemoryPipeline(self, transaction=transaction) class InMemoryPipeline: + """A fake Redis pipeline that records any method call and replays it later.""" + def __init__(self, backend: InMemoryCacheBackend, *, transaction: bool) -> None: + """Queue commands for ``backend``; ``transaction`` makes ``execute`` atomic.""" self.backend = backend self.transaction = transaction self.operations: list[tuple[str, tuple[object, ...], dict[str, object]]] = [] def __getattr__(self, name: str) -> Callable[..., InMemoryPipeline]: + """Return a recorder that queues a call to the backend method ``name``.""" + def queue(*args: object, **kwargs: object) -> InMemoryPipeline: self.operations.append((name, args, kwargs)) return self @@ -149,6 +184,11 @@ def queue(*args: object, **kwargs: object) -> InMemoryPipeline: return queue def execute(self) -> list[object]: + """Replay the queued commands on the backend, then clear the queue. + + With ``transaction`` the replay holds the backend lock, as MULTI/EXEC + would; without it, other callers' commands may interleave. + """ lock = self.backend._lock if self.transaction else RLock() with lock: results = [ @@ -159,9 +199,11 @@ def execute(self) -> list[object]: return results def __enter__(self) -> Self: + """Return the pipeline itself for use in a ``with`` block.""" return self def __exit__(self, *_args: object) -> None: + """Drop queued commands that were never executed.""" self.operations.clear() @@ -169,27 +211,35 @@ class DisabledCacheBackend(InMemoryCacheBackend): """No-storage backend for unselected unit-test/application cache mode.""" def get(self, key: str) -> None: - return None + """Return None: every read is a miss.""" + return None # noqa: RET501 # explicit None mirrors the Redis get contract def set(self, key: str, value: str, **kwargs: object) -> bool: + """Discard the value and report the write as not stored.""" return False def delete(self, *keys: str) -> int: + """Return 0: no key ever exists.""" return 0 def zadd(self, key: str, mapping: dict[str, float]) -> int: + """Discard the members and report none added.""" return 0 def zrevrange(self, key: str, start: int, end: int) -> list[str]: + """Return an empty list: no sorted set ever exists.""" return [] def zrem(self, key: str, *members: str) -> int: + """Return 0: there is nothing to remove.""" return 0 def zremrangebyrank(self, key: str, start: int, end: int) -> int: + """Return 0: there is nothing to remove.""" return 0 def expire(self, key: str, seconds: int) -> bool: + """Return False, as Redis does for a key that does not exist.""" return False def eval( @@ -198,4 +248,5 @@ def eval( numkeys: int, *keys_and_args: str, ) -> int: + """Raise RuntimeError, so claim operations fail closed when disabled.""" raise RuntimeError("the shared runtime cache is disabled") diff --git a/policyengine_api/runtime_cache/household_calculations.py b/policyengine_api/runtime_cache/household_calculations.py index f490c3bd2..1082054ef 100644 --- a/policyengine_api/runtime_cache/household_calculations.py +++ b/policyengine_api/runtime_cache/household_calculations.py @@ -21,6 +21,13 @@ @dataclass(frozen=True) class HouseholdCalculationIdentity: + """Every input that determines a calculated household; the cache key hashes it. + + The content hashes and package versions make edited inputs and upgraded + models miss rather than reuse stale results. ``spm`` is the requested + resolved SPM configuration, or None when the request selected no SPM. + """ + country_id: str household_id: int policy_id: int @@ -33,6 +40,8 @@ class HouseholdCalculationIdentity: @dataclass(frozen=True) class CachedHouseholdCalculation: + """A calculated household, its warnings, and the SPM receipt if SPM was used.""" + household: dict[str, Any] warnings: tuple[str, ...] = () spm_provenance: dict[str, Any] | None = None @@ -42,6 +51,7 @@ class HouseholdCalculationCache: """Cache a calculated household, its warnings and its receipts atomically.""" def __init__(self, client: CacheBackend, namespace: CacheNamespace) -> None: + """Store results in ``client`` under ``namespace`` for about a day.""" self._cache = RecoverableJSONCache( client, namespace, @@ -51,12 +61,18 @@ def __init__(self, client: CacheBackend, namespace: CacheNamespace) -> None: ) def cache_key(self, identity: HouseholdCalculationIdentity) -> str: + """Return the backend key a result for ``identity`` is stored under.""" return self._cache.key(asdict(identity)) def get( self, identity: HouseholdCalculationIdentity, ) -> CachedHouseholdCalculation | None: + """Return the cached result, or None on a miss or any malformed entry. + + An entry whose SPM receipt is missing, invalid, or resolves to a + different SPM configuration than ``identity`` asks for is also a miss. + """ payload = self._cache.get(asdict(identity)) if not isinstance(payload, dict): return None @@ -96,6 +112,11 @@ def set( identity: HouseholdCalculationIdentity, value: CachedHouseholdCalculation, ) -> bool: + """Cache a result, returning False if it was not written. + + Results whose SPM receipt does not match the identity's SPM request are + refused rather than stored, so a later read cannot serve them. + """ if identity.spm is None: if value.spm_provenance is not None: return False diff --git a/policyengine_api/runtime_cache/reform_impacts.py b/policyengine_api/runtime_cache/reform_impacts.py index 9012010b3..32d7a867d 100644 --- a/policyengine_api/runtime_cache/reform_impacts.py +++ b/policyengine_api/runtime_cache/reform_impacts.py @@ -34,6 +34,7 @@ class ReformImpactStartClaim: observability_id: str def to_payload(self) -> dict[str, str]: + """Return the JSON object stored as the claim's value.""" return { "submission_claim_id": self.submission_claim_id, "observability_id": self.observability_id, @@ -41,6 +42,7 @@ def to_payload(self) -> dict[str, str]: @classmethod def from_payload(cls, payload: object) -> "ReformImpactStartClaim | None": + """Rebuild a claim, or return None if a field is missing or invalid.""" if not isinstance(payload, dict): return None submission_claim_id = payload.get("submission_claim_id") @@ -57,6 +59,11 @@ def from_payload(cls, payload: object) -> "ReformImpactStartClaim | None": @dataclass(frozen=True) class CachedReformImpact: + """An economy-impact job and its result, shaped like a legacy reform_impact row. + + ``status`` is ``computing``, ``ok`` or ``error``; naive datetimes are UTC. + """ + reform_impact_id: int baseline_policy_id: int reform_policy_id: int @@ -162,6 +169,7 @@ class ReformImpactCache: family = "reform-impact" def __init__(self, client: CacheBackend, namespace: CacheNamespace) -> None: + """Keep values, indexes and start claims in ``client`` under ``namespace``.""" self.client = client self.namespace = namespace self._start_claims = ExpiringClaimStore(client, family=self.family) @@ -397,6 +405,11 @@ def set( *, record_execution_id: str | None = None, ) -> bool: + """Store an impact and index it by scope and recency in one transaction. + + Each index keeps only the newest 1,000 entries. Return False if the + write failed; raise ValueError when there is no record identifier. + """ # Terminal errors clear the worker handle in the payload while keeping # the original record key and lookup indexes stable. record_execution_id = record_execution_id or impact.execution_id @@ -436,6 +449,7 @@ def set( return True def get_by_execution_id(self, execution_id: str) -> CachedReformImpact | None: + """Return the job's impact, or None if absent, unreadable or unreachable.""" try: value = self.client.get(self._record_key(execution_id)) except Exception: @@ -467,6 +481,7 @@ def _from_index(self, index: str, limit: int) -> list[CachedReformImpact]: return [impact for impact in impacts if impact is not None] def recent(self, limit: int) -> list[CachedReformImpact]: + """Return up to ``limit`` impacts, most recently started first.""" return self._from_index(self._recent_index(), limit) def matching( @@ -482,6 +497,12 @@ def matching( options_hash: str, options_hash_pattern: str | None = None, ) -> list[CachedReformImpact]: + """Return the cached impacts for one policy pair and scope, best match first. + + Exact ``options_hash`` matches sort before those matching + ``options_hash_pattern``, a SQL ``LIKE`` pattern; ties go to the most + recently started. Only the scope index's newest 1,000 entries are read. + """ probe = CachedReformImpact( reform_impact_id=0, baseline_policy_id=baseline_policy_id, @@ -531,6 +552,11 @@ def update( clear_execution_id: bool = False, **changes: Any, # noqa: ANN401 # forwarded unchanged to dataclasses.replace ) -> CachedReformImpact | None: + """Apply field changes to a cached impact and rewrite it under the same key. + + ``clear_execution_id`` drops the worker handle without moving the record. + Return the updated impact, or None if it is missing or the write failed. + """ impact = self.get_by_execution_id(execution_id) if impact is None: return None @@ -550,6 +576,10 @@ def delete_matching_computing( time_period: str, options_hash: str, ) -> None: + """Remove in-flight (``computing``) impacts for this scope and options. + + Cache failures are ignored, so a stale entry may survive until it expires. + """ # Deletion historically omits api_version, so search the bounded recent # index and remove only matching in-flight values. impacts = self.recent(REFORM_IMPACT_INDEX_LIMIT) diff --git a/policyengine_api/runtime_cache/settings.py b/policyengine_api/runtime_cache/settings.py index d7b293d0f..958b0ac43 100644 --- a/policyengine_api/runtime_cache/settings.py +++ b/policyengine_api/runtime_cache/settings.py @@ -33,6 +33,8 @@ class RuntimeCacheConfigurationError(RuntimeError): @dataclass(frozen=True) class RuntimeCacheSettings: + """Validated cache configuration; the URL and CA bundle stay secret in reprs.""" + mode: str environment: str service: str @@ -45,6 +47,7 @@ class RuntimeCacheSettings: @property def enabled(self) -> bool: + """Whether a real Redis backend is configured (local or deployed mode).""" return self.mode != "disabled" From a455332a505e7198a03faf0ea3618dabdafd5b43 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:53:03 +0400 Subject: [PATCH 04/12] Document the rules tool and repository scripts Add docstrings to scripts/rules (the registry model, loader, resolver, validator and renderer), the quality-guard runner and migration guard, the migration baseline, mirroring and smoke-test scripts, the Alembic environments, and the .github helper scripts. Reword the rules tool's noun-phrase function summaries into the imperative and give its multi-line summaries a separating blank line. scripts/rules is the canonical copy of a tool shared byte-for-byte with policyengine-sim-api, and the argparse help taken from module docstrings is unchanged. Co-Authored-By: Claude Opus 5.5 --- .github/bump_version.py | 9 ++++++ .github/find-api-model-versions.py | 6 ++++ .github/scripts/check_updates.py | 6 ++++ .../scripts/verify_observability_collector.py | 7 +++++ dashboard/experiments/gpt4_api_user.py | 15 +++++++++ migrations/v1/env.py | 2 ++ migrations/v2/env.py | 6 ++++ scripts/capture_migration_baseline.py | 17 ++++++++++ scripts/compare_migration_baseline.py | 6 ++++ scripts/export_migration_contracts.py | 1 + scripts/guards/migration_contracts.py | 2 ++ scripts/measure_cloud_run_runtime.py | 1 + scripts/process_v1_household_mirror_event.py | 5 +++ scripts/rules/adapters.py | 4 +++ scripts/rules/baseline.py | 6 +++- scripts/rules/ci_scripts.py | 19 ++++++++++-- scripts/rules/cli.py | 4 +++ scripts/rules/discover.py | 1 + scripts/rules/guard.py | 7 +++++ scripts/rules/load.py | 2 ++ scripts/rules/model.py | 30 ++++++++++++++++++ scripts/rules/render.py | 10 +++++- scripts/rules/resolve.py | 31 +++++++++++++------ scripts/rules/validate.py | 5 +++ scripts/rules/workflows.py | 3 ++ scripts/run_quality_guards.py | 1 + scripts/smoke_test_local_boot.py | 8 +++++ scripts/v1_database_migration.py | 26 ++++++++++++++++ 28 files changed, 226 insertions(+), 14 deletions(-) diff --git a/.github/bump_version.py b/.github/bump_version.py index 1e7b49e18..e29c6bf8c 100644 --- a/.github/bump_version.py +++ b/.github/bump_version.py @@ -6,6 +6,7 @@ def get_current_version(pyproject_path: Path) -> str: + """Return the version in pyproject.toml, exiting with status 1 if it is absent.""" text = pyproject_path.read_text() match = re.search(r'^version\s*=\s*"(\d+\.\d+\.\d+)"', text, re.MULTILINE) if not match: @@ -18,6 +19,11 @@ def get_current_version(pyproject_path: Path) -> str: def infer_bump(changelog_dir: Path) -> str: + """Return "major", "minor" or "patch" from the pending changelog fragment types. + + Any breaking fragment makes a major bump and any added or removed fragment a + minor one; exit with status 1 when there are no fragments. + """ fragments = [ f for f in changelog_dir.iterdir() if f.is_file() and f.name != ".gitkeep" ] @@ -39,6 +45,7 @@ def infer_bump(changelog_dir: Path) -> str: def bump_version(version: str, bump: str) -> str: + """Return ``version`` raised by one ``bump`` level, resetting the lower parts.""" major, minor, patch = (int(x) for x in version.split(".")) if bump == "major": return f"{major + 1}.0.0" @@ -48,6 +55,7 @@ def bump_version(version: str, bump: str) -> str: def update_file(path: Path, old_version: str, new_version: str) -> None: + """Rewrite the ``version = "..."`` line in ``path``; leave other files untouched.""" text = path.read_text() updated = text.replace( f'version = "{old_version}"', @@ -59,6 +67,7 @@ def update_file(path: Path, old_version: str, new_version: str) -> None: def main() -> None: + """Bump the version in pyproject.toml according to the changelog.d fragments.""" root = Path(__file__).resolve().parent.parent pyproject = root / "pyproject.toml" changelog_dir = root / "changelog.d" diff --git a/.github/find-api-model-versions.py b/.github/find-api-model-versions.py index 55287ca08..6f80d3f57 100644 --- a/.github/find-api-model-versions.py +++ b/.github/find-api-model-versions.py @@ -1,3 +1,9 @@ +"""CI helper that reads the versions pinned by the installed PolicyEngine bundle. + +It checks the package and data versions against the API's version constants, +then writes them to GITHUB_ENV or, with --shell, prints shell assignments. +""" + import argparse import os import shlex diff --git a/.github/scripts/check_updates.py b/.github/scripts/check_updates.py index 4d59de47b..e3c232e9f 100644 --- a/.github/scripts/check_updates.py +++ b/.github/scripts/check_updates.py @@ -24,6 +24,8 @@ class ChangelogEntry(TypedDict): + """One released version's changelog items, grouped by lowercase category.""" + version: str changes: dict[str, list[str]] @@ -118,6 +120,10 @@ def format_changes(entries: list[ChangelogEntry]) -> str: def main() -> int: + """Print the markdown changes between two versions of a package. + + Always return 0, so an unreachable changelog never fails the update job. + """ parser = argparse.ArgumentParser( description="Fetch changelog entries between two package versions." ) diff --git a/.github/scripts/verify_observability_collector.py b/.github/scripts/verify_observability_collector.py index 7d387a93c..633c7213f 100644 --- a/.github/scripts/verify_observability_collector.py +++ b/.github/scripts/verify_observability_collector.py @@ -64,6 +64,8 @@ class VerificationProbe(NamedTuple): + """Identifiers of one verification export, for finding its stored telemetry.""" + verification_id: str trace_id: str started_at: float @@ -369,6 +371,11 @@ def verify( def main() -> None: + """Verify the collector, reading both tokens from the environment. + + COLLECTOR_ID_TOKEN authenticates to the collector and GOOGLE_OAUTH_ACCESS_TOKEN + reads Cloud Trace and Cloud Monitoring. + """ parser = argparse.ArgumentParser() parser.add_argument("--endpoint", required=True) parser.add_argument("--project-id", required=True) diff --git a/dashboard/experiments/gpt4_api_user.py b/dashboard/experiments/gpt4_api_user.py index 27b21f890..1fd37047b 100644 --- a/dashboard/experiments/gpt4_api_user.py +++ b/dashboard/experiments/gpt4_api_user.py @@ -1,3 +1,10 @@ +"""Streamlit experiment that answers a UK policy question with GPT-4 and the API. + +It finds related parameters by embedding similarity, asks GPT-4 to write the +reform as JSON, runs it on the public PolicyEngine API, and asks GPT-4 to answer +from the resulting economic impact. +""" + import json import os import time @@ -14,6 +21,7 @@ @st.cache_data def get_embeddings_df() -> pd.DataFrame: + """Load parameter_embeddings.csv.gz, evaluating each embedding string to a list.""" embeddings_df = pd.read_csv("parameter_embeddings.csv.gz", compression="gzip") embeddings_df.parameter_embedding = embeddings_df.parameter_embedding.apply( lambda x: eval(x) @@ -23,6 +31,7 @@ def get_embeddings_df() -> pd.DataFrame: @st.cache_data def embed(prompt: str, engine: str = "text-embedding-ada-002") -> list[float]: + """Return the OpenAI embedding of ``prompt``, cached across Streamlit reruns.""" return get_embedding(prompt, engine=engine) @@ -59,6 +68,7 @@ def embed(prompt: str, engine: str = "text-embedding-ada-002") -> list[float]: def write_prompt(question: str, relevant_parameters: object) -> str: + """Build the GPT-4 prompt that rewrites ``question`` as a reform JSON object.""" return f"""Question: {question} Here's some metadata about the parameters that are relevant to your question: Relevant parameters: {relevant_parameters} @@ -79,6 +89,7 @@ def write_prompt(question: str, relevant_parameters: object) -> str: @st.cache_data def get_gpt4_response(prompt: str) -> str: + """Return GPT-4's reply to ``prompt`` as one user message, cached across reruns.""" return openai.ChatCompletion.create( model="gpt-4", messages=[{"role": "user", "content": prompt}], @@ -94,6 +105,10 @@ def get_gpt4_response(prompt: str) -> str: @st.cache_data def get_policyengine_impact(reform: object) -> object: + """Create ``reform`` on the public UK API and return its 2023 UK economic impact. + + Poll every five seconds, with no time limit, while the impact is computing. + """ payload = {"data": reform} url = "https://api.policyengine.org/uk/policy" diff --git a/migrations/v1/env.py b/migrations/v1/env.py index ab4b6983e..c2ce10a5c 100644 --- a/migrations/v1/env.py +++ b/migrations/v1/env.py @@ -40,6 +40,7 @@ def run_migrations_offline() -> None: + """Emit the migration SQL for the configured MySQL URL without connecting.""" context.configure( url=config.get_main_option("sqlalchemy.url"), target_metadata=target_metadata, @@ -52,6 +53,7 @@ def run_migrations_offline() -> None: def run_migrations_online() -> None: + """Migrate the connection the caller supplied, or else a new configured engine.""" if provided_connection is not None: _run_migrations_with_connection(provided_connection) return diff --git a/migrations/v2/env.py b/migrations/v2/env.py index cbde16f98..302ca39c5 100644 --- a/migrations/v2/env.py +++ b/migrations/v2/env.py @@ -64,6 +64,7 @@ def _configure(connection: Connection) -> None: def run_migrations_offline() -> None: + """Raise, because v2 migrations need a live connection to qualify the target.""" raise RuntimeError( "v2 migrations require an online connection so target identity and " "generated application-data before/after states can be verified" @@ -71,6 +72,11 @@ def run_migrations_offline() -> None: def run_migrations_online() -> None: + """Migrate the caller's connection, or else a new engine in one transaction. + + The new engine's transaction commits, so the target qualification's reads + cannot leave the completed migration to be rolled back. + """ provided_connection = config.attributes.get("connection") if provided_connection is not None: _configure(provided_connection) diff --git a/scripts/capture_migration_baseline.py b/scripts/capture_migration_baseline.py index 75e37dcd0..674871fcf 100644 --- a/scripts/capture_migration_baseline.py +++ b/scripts/capture_migration_baseline.py @@ -24,6 +24,8 @@ @dataclass(frozen=True) class RouteProbe: + """One API request to time, and the status codes that count as success.""" + name: str method: str path: str @@ -32,6 +34,8 @@ class RouteProbe: @dataclass(frozen=True) class ProbeResult: + """One timed request, or, when ``completed`` is set, one whole simulation run.""" + name: str method: str path: str @@ -105,6 +109,7 @@ def _request_probe( def run_probes(base_url: str, repetitions: int) -> list[ProbeResult]: + """Time each default probe ``repetitions`` times; transport errors are failures.""" results: list[ProbeResult] = [] with httpx.Client(base_url=base_url.rstrip("/"), timeout=90.0) as client: for _ in range(repetitions): @@ -127,6 +132,12 @@ def run_simulation_gateway_probe( poll_timeout_seconds: float, poll_interval_seconds: float, ) -> list[ProbeResult]: + """Submit one economy comparison to the gateway and poll it until it finishes. + + Return each submit and poll request followed by one completion row, whose + ``completed`` flag is False on failure or when the timeout (counted from + submission) expires. + """ results: list[ProbeResult] = [] started_at = time.perf_counter() @@ -220,6 +231,11 @@ def run_simulation_gateway_probe( def summarize(results: Iterable[ProbeResult]) -> dict: + """Summarise request latency, error rate and simulation completion times. + + Completion rows are counted apart from request rows, and latency + percentiles use successful requests only. + """ rows = list(results) request_rows = [row for row in rows if row.completed is None] latencies = [row.latency_ms for row in request_rows if row.ok] @@ -255,6 +271,7 @@ def summarize(results: Iterable[ProbeResult]) -> dict: def main(argv: list[str] | None = None) -> int: + """Print a JSON baseline summary; return 0 even when probes fail or are skipped.""" parser = argparse.ArgumentParser() parser.add_argument("--base-url", default=os.environ.get("API_BASE_URL")) parser.add_argument( diff --git a/scripts/compare_migration_baseline.py b/scripts/compare_migration_baseline.py index affda18a5..d3f320526 100644 --- a/scripts/compare_migration_baseline.py +++ b/scripts/compare_migration_baseline.py @@ -31,6 +31,11 @@ @dataclass(frozen=True) class ProbeStats: + """Error rate and p95 of one probe, or of the simulation completion pseudo-probe. + + For the completion pseudo-probe ``p95_latency_ms`` holds completion time. + """ + name: str request_count: int failure_count: int @@ -138,6 +143,7 @@ def compare( def main(argv: list[str] | None = None) -> int: + """Print the per-probe gate report; return 1 if any baseline probe fails.""" parser = argparse.ArgumentParser( description="Gate a candidate baseline JSON against a reference " "baseline JSON, per probe." diff --git a/scripts/export_migration_contracts.py b/scripts/export_migration_contracts.py index 04638a9eb..1a716dc65 100644 --- a/scripts/export_migration_contracts.py +++ b/scripts/export_migration_contracts.py @@ -146,6 +146,7 @@ def write_outputs( def main(argv: list[str] | None = None) -> int: + """Regenerate the migration contract JSON and Markdown, overwriting both files.""" parser = argparse.ArgumentParser() parser.add_argument("--json", default=str(DEFAULT_JSON)) parser.add_argument("--markdown", default=str(DEFAULT_MARKDOWN)) diff --git a/scripts/guards/migration_contracts.py b/scripts/guards/migration_contracts.py index 2e61125c5..a7760ea43 100644 --- a/scripts/guards/migration_contracts.py +++ b/scripts/guards/migration_contracts.py @@ -121,6 +121,7 @@ def _check_generated_artifacts(payload: dict[str, Any]) -> list[str]: def check() -> list[str]: + """Return every route-group, workflow and stale-artifact violation.""" payload = export_migration_contracts.build_payload() return [ *_check_route_groups(payload), @@ -130,6 +131,7 @@ def check() -> list[str]: def main() -> int: + """Print the guard's violations and return 1 if there are any.""" violations = check() if not violations: print("migration-contracts guard passed") diff --git a/scripts/measure_cloud_run_runtime.py b/scripts/measure_cloud_run_runtime.py index e269315b0..3d1c049d3 100644 --- a/scripts/measure_cloud_run_runtime.py +++ b/scripts/measure_cloud_run_runtime.py @@ -233,6 +233,7 @@ async def _run(args: argparse.Namespace) -> dict: def main() -> int: + """Wait for readiness, run the load cell, and print its JSON summary to stdout.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--base-url", required=True) parser.add_argument("--cell-label", required=True) diff --git a/scripts/process_v1_household_mirror_event.py b/scripts/process_v1_household_mirror_event.py index c4e5d3c0b..61908eba0 100644 --- a/scripts/process_v1_household_mirror_event.py +++ b/scripts/process_v1_household_mirror_event.py @@ -194,6 +194,11 @@ def main( *, environ: Mapping[str, str] | None = None, ) -> int: + """Mirror the selected pending event and print a JSON outcome; return 1 on failure. + + Error output names only configuration problems; any other failure is reported + by exception type, so source data and credentials never reach the log. + """ args = _parser().parse_args(argv) if args.legacy_household_id < 0: _parser().error("--legacy-household-id must be non-negative") diff --git a/scripts/rules/adapters.py b/scripts/rules/adapters.py index b22bd112a..19a7ee979 100644 --- a/scripts/rules/adapters.py +++ b/scripts/rules/adapters.py @@ -18,6 +18,10 @@ def check_adapters(agents: Agents, root: Path) -> list[Violation]: + """Return a violation for each way an adapter stops being a short pointer. + + Missing adapter files are skipped rather than reported. + """ violations = [] instructions = Path(agents.instructions).name for adapter in agents.adapters: diff --git a/scripts/rules/baseline.py b/scripts/rules/baseline.py index f1a8aeb4b..5bc3feb90 100644 --- a/scripts/rules/baseline.py +++ b/scripts/rules/baseline.py @@ -18,12 +18,16 @@ class BaseUnavailable(Exception): def default_base_ref() -> str | None: + """Return ``origin/`` on a pull-request run, otherwise None.""" base = os.environ.get("GITHUB_BASE_REF") return f"origin/{base}" if base else None def load_base_registry(root: Path, ref: str) -> Registry | None: - """The registry as it is on ``ref``, or None if the base has no registry yet.""" + """Return the registry as it is on ``ref``, or None if the base has none yet. + + Raise BaseUnavailable when ``ref`` does not name a commit. + """ verify = subprocess.run( ["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"], cwd=root, diff --git a/scripts/rules/ci_scripts.py b/scripts/rules/ci_scripts.py index 071b2560f..0a0edecd2 100644 --- a/scripts/rules/ci_scripts.py +++ b/scripts/rules/ci_scripts.py @@ -36,6 +36,8 @@ @dataclass class Findings: + """What one ci-scripts run found, as baseline-style entries for each check.""" + multiline: set[str] = field(default_factory=set) untested: set[str] = field(default_factory=set) shellcheck: set[str] | None = None # None when shellcheck could not run @@ -43,8 +45,10 @@ class Findings: def multiline_runs(config: CiScripts, root: Path) -> set[str]: - """``workflow :: job :: step`` for every run: block with more than one line of - code, or with commands chained by ``&&``, ``||`` or ``;``. + """Return ``workflow :: job :: step`` for each run: block that is not one command. + + A block fails when it has more than one line of code, or chains commands with + ``&&``, ``||`` or ``;``. """ found = set() for workflow in sorted(root.glob(config.workflows)): @@ -67,7 +71,7 @@ def called_scripts(config: CiScripts, root: Path) -> set[str]: def untested_scripts(config: CiScripts, root: Path) -> set[str]: - """Called scripts whose file name appears in no test file.""" + """Return the called scripts whose file name appears in no test file.""" tests = [ path.read_text(errors="replace") for pattern in config.test_globs @@ -102,6 +106,7 @@ def shellcheck_failures(config: CiScripts, root: Path) -> set[str] | None: def find(config: CiScripts, root: Path) -> Findings: + """Run the three checks on freshly parsed workflows, noting a skipped shellcheck.""" clear_caches() findings = Findings( multiline=multiline_runs(config, root), @@ -139,6 +144,10 @@ def _rule_for(registry: Registry, baseline: str, check: str) -> str | None: def compare(registry: Registry, findings: Findings, root: Path) -> list[Violation]: + """Report findings missing from their baseline and baseline entries now fixed. + + A check that could not run is skipped rather than reported. + """ config = registry.ci_scripts assert config is not None checks = ( @@ -205,6 +214,10 @@ def write_baselines(registry: Registry, findings: Findings, root: Path) -> None: def run_check(root: Path = REPO_ROOT) -> tuple[list[Violation], list[str]]: + """Load the registry and return the ci-scripts violations and notes. + + A missing shellcheck is a violation on GitHub Actions and only a note elsewhere. + """ registry, violations = load_registry(root) if registry is None or registry.ci_scripts is None: return violations, [ diff --git a/scripts/rules/cli.py b/scripts/rules/cli.py index 9322bb086..a4134cacb 100644 --- a/scripts/rules/cli.py +++ b/scripts/rules/cli.py @@ -24,6 +24,10 @@ def _print(violations: list, notes: list[str] | None = None) -> int: def main(argv: list[str] | None = None) -> int: + """Run one rules command and return 1 if it found violations. + + ``render`` refuses to write anything while the registry fails validation. + """ parser = argparse.ArgumentParser( prog="python -m scripts.rules", description=__doc__ ) diff --git a/scripts/rules/discover.py b/scripts/rules/discover.py index 323425618..a352c05ef 100644 --- a/scripts/rules/discover.py +++ b/scripts/rules/discover.py @@ -60,6 +60,7 @@ def _discovered_tests(registry: Registry, root: Path) -> set[str]: def discover_hidden(registry: Registry, root: Path) -> list[Violation]: + """Return a hidden-rule violation for each check no rule references or ignores.""" violations = [] rules = registry.rules diff --git a/scripts/rules/guard.py b/scripts/rules/guard.py index c68fce02b..3efa3d4f5 100644 --- a/scripts/rules/guard.py +++ b/scripts/rules/guard.py @@ -38,6 +38,8 @@ @dataclass class GuardResult: + """Guard violations, which fail the run, and notes, which only inform.""" + violations: list[Violation] = field(default_factory=list) notes: list[str] = field(default_factory=list) @@ -75,6 +77,11 @@ def run_guard( today: date | None = None, compare_base: bool = True, ) -> GuardResult: + """Run every registry check, then compare baselines with the base ref. + + Stop after the schema errors if the registry does not load. An unreadable + base ref is a violation on GitHub Actions and only a note elsewhere. + """ clear_caches() result = GuardResult() registry, schema_violations = load_registry(root) diff --git a/scripts/rules/load.py b/scripts/rules/load.py index 83b455e2e..936a3a9f4 100644 --- a/scripts/rules/load.py +++ b/scripts/rules/load.py @@ -389,6 +389,7 @@ def parse(raw: dict[str, Any]) -> tuple[Registry | None, list[Violation]]: def load_text(text: str) -> tuple[Registry | None, list[Violation]]: + """Parse registry TOML text, or return None and one violation if it is not TOML.""" try: raw = tomllib.loads(text) except tomllib.TOMLDecodeError as error: @@ -404,6 +405,7 @@ def load_text(text: str) -> tuple[Registry | None, list[Violation]]: def load_registry(root: Path) -> tuple[Registry | None, list[Violation]]: + """Load the registry under ``root``, or return None and a violation if missing.""" path = root / REGISTRY_PATH if not path.exists(): return None, [ diff --git a/scripts/rules/model.py b/scripts/rules/model.py index 6d8f9c365..401becaaa 100644 --- a/scripts/rules/model.py +++ b/scripts/rules/model.py @@ -12,12 +12,16 @@ class Tier(StrEnum): + """How mechanically a rule can be checked; judgment-tier rules never block.""" + MECHANICAL = "mechanical" METRIC = "metric" JUDGMENT = "judgment" class Status(StrEnum): + """How a rule is enforced today; only enforced and ratcheted rules can block.""" + ENFORCED = "enforced" RATCHETED = "ratcheted" ADVISORY = "advisory" @@ -25,6 +29,8 @@ class Status(StrEnum): class Kind(StrEnum): + """The checker behind a rule; it decides which ``[rule.enforcer]`` fields apply.""" + PYTEST = "pytest" QUALITY_GUARD = "quality-guard" RUFF = "ruff" @@ -47,12 +53,15 @@ class Violation: fix: str def __str__(self) -> str: + """Format as ``[rule] check: message. Fix: fix.``; rule defaults to registry.""" subject = self.rule_id or "registry" return f"[{subject}] {self.check}: {self.message}. Fix: {self.fix}." @dataclass(frozen=True) class Category: + """A ``[[category]]`` grouping rules, in declared order, in the rendered views.""" + id: str title: str @@ -67,6 +76,8 @@ class Ignored: @dataclass(frozen=True) class Discovery: + """Where the guard looks for unregistered checks, and which ones it may skip.""" + workflows: tuple[str, ...] = (".github/workflows/pr.yml",) guard_test_globs: tuple[str, ...] = () extra_guard_tests: tuple[str, ...] = () @@ -97,6 +108,8 @@ class CiScripts: @dataclass(frozen=True) class Enforcer: + """What checks a rule: its ``kind`` plus the raw ``[rule.enforcer]`` fields.""" + kind: Kind fields: Mapping[str, Any] = field(default_factory=dict) @@ -105,6 +118,7 @@ def get( key: str, default: object = None, ) -> Any: # noqa: ANN401 # a raw TOML value that callers read by its shape + """Return a raw enforcer field, or ``default`` when the registry omits it.""" return self.fields.get(key, default) @property @@ -153,6 +167,8 @@ def expects_proof(self) -> bool: @dataclass(frozen=True) class CiRef: + """A ``[[rule.ci]]`` reference: a workflow job whose run text has ``invocation``.""" + workflow: str job: str invocation: str @@ -160,17 +176,26 @@ class CiRef: @dataclass(frozen=True) class Proof: + """The raw ``[rule.proof]`` fields; which keys are required depends on the tier.""" + fields: Mapping[str, Any] @dataclass(frozen=True) class Baseline: + """A ratcheted rule's file of tolerated violations and the number it lists. + + The count must match the file and may never grow relative to the base branch. + """ + file: str count: int @dataclass(frozen=True) class RuleException: + """A temporary carve-out of ``path`` from a rule; the guard fails once expired.""" + path: str reason: str expires: date @@ -178,6 +203,8 @@ class RuleException: @dataclass(frozen=True) class Rule: + """One ``[[rule]]`` registry entry; rules-registry.md explains each field.""" + id: str category: str statement: str @@ -205,6 +232,8 @@ def blocking(self) -> bool: @dataclass(frozen=True) class Registry: + """The parsed contents of rules/registry.toml.""" + schema_version: int repo: str categories: tuple[Category, ...] @@ -214,6 +243,7 @@ class Registry: ci_scripts: CiScripts | None = None def rule(self, rule_id: str) -> Rule | None: + """Return the rule with ``rule_id``, or None if the registry has none.""" for rule in self.rules: if rule.id == rule_id: return rule diff --git a/scripts/rules/render.py b/scripts/rules/render.py index 97815bf30..2b70361e4 100644 --- a/scripts/rules/render.py +++ b/scripts/rules/render.py @@ -52,7 +52,7 @@ def _proof_label(rule: Rule) -> str: def _enforced_by(rule: Rule) -> str: - """What enforces a rule, naming the CI job when the job itself is the check.""" + """Describe what enforces a rule, naming the CI job when the job is the check.""" if rule.enforcer.kind is Kind.WORKFLOW_JOB: jobs = ", ".join(f"`{ref.job}`" for ref in rule.ci) return f"CI job {jobs}" @@ -121,6 +121,7 @@ def _bullets(items: list[str]) -> list[str]: def render_rules_md(registry: Registry) -> str: + """Render rules.md: the tier summary, the gaps, then each category's rules.""" rules = sorted(registry.rules, key=lambda r: r.id) out = [ "# Rules", @@ -211,6 +212,11 @@ def _agents_line(rule: Rule) -> str: def render_agents_block(registry: Registry) -> str: + """Render the AGENTS.md rules block, markers included. + + Past ``AGENTS_COLLAPSE_AT`` rules, each category lists its checked rules on + one line and spells out only the doc-only ones. + """ out = [ BEGIN, "", @@ -245,6 +251,7 @@ def splice_block(text: str, block: str) -> str: def extract_block(text: str) -> str | None: + """Return the generated rules block in ``text``, markers included, or None.""" start, end = text.find(BEGIN), text.find(END) if start == -1 or end == -1 or end < start: return None @@ -252,6 +259,7 @@ def extract_block(text: str) -> str | None: def write_outputs(registry: Registry, root: Path) -> list[Path]: + """Write rules.md and splice the rules block into AGENTS.md; return both paths.""" rules_md = root / RULES_MD rules_md.parent.mkdir(parents=True, exist_ok=True) rules_md.write_text(render_rules_md(registry)) diff --git a/scripts/rules/resolve.py b/scripts/rules/resolve.py index cddb6536d..ed2bba559 100644 --- a/scripts/rules/resolve.py +++ b/scripts/rules/resolve.py @@ -34,6 +34,7 @@ def _read(path: Path) -> str | None: def clear_caches() -> None: + """Forget cached file reads and parses, so the next check sees the files afresh.""" _read.cache_clear() workflow_jobs.cache_clear() test_names.cache_clear() @@ -41,6 +42,7 @@ def clear_caches() -> None: @cache def workflow_jobs(path: Path) -> tuple[dict[str, JobBlock], tuple[str, ...]]: + """Return a workflow's jobs by id and its parse problems; cached per path.""" text = _read(path) if text is None: return {}, (f"{path.name} is missing",) @@ -97,7 +99,10 @@ def node_exists(root: Path, node: str) -> str | None: def guard_names(root: Path, runner: str = DEFAULT_RUNNER) -> list[str] | None: - """String names in the ``GUARDS`` tuple of the quality-guard runner.""" + """Return the string names in the runner's ``GUARDS`` tuple. + + Return None if the runner is missing, and an empty list if it has no ``GUARDS``. + """ text = _read(root / runner) if text is None: return None @@ -131,7 +136,7 @@ def _toml(path: Path) -> dict[str, Any]: def ruff_lint_config(root: Path) -> dict[str, Any]: - """The ``lint`` table from ruff.toml, or ``[tool.ruff.lint]`` in pyproject.""" + """Return ruff.toml's ``lint`` table, else ``[tool.ruff.lint]`` in pyproject.""" ruff_toml = _toml(root / "ruff.toml") if ruff_toml: return ruff_toml.get("lint", {}) @@ -141,12 +146,14 @@ def ruff_lint_config(root: Path) -> dict[str, Any]: def ruff_selected(root: Path) -> tuple[str, ...]: + """Return ruff's ``select`` (its default when unset) plus ``extend-select``.""" lint = ruff_lint_config(root) selected = tuple(lint.get("select", RUFF_DEFAULT_SELECT)) return selected + tuple(lint.get("extend-select", ())) def mypy_files(root: Path) -> list[str]: + """Return ``[tool.mypy].files`` as a list, even when it is a single string.""" files = ( _toml(root / "pyproject.toml").get("tool", {}).get("mypy", {}).get("files", []) ) @@ -323,6 +330,7 @@ def _resolve_workflow_job(rule: Rule, root: Path) -> list[Violation]: def resolve_enforcer(rule: Rule, root: Path) -> list[Violation]: + """Check that the rule's enforcer exists where its kind says it should.""" return RESOLVERS[rule.enforcer.kind](rule, root) @@ -332,7 +340,7 @@ def resolve_enforcer(rule: Rule, root: Path) -> list[Violation]: def _job_for( rule_id: str, ref: CiRef, root: Path, check: str, field: str ) -> tuple[JobBlock | None, list[Violation]]: - """The job a CI reference names, if it exists and runs the invocation.""" + """Return the job a CI reference names if it runs the invocation, plus problems.""" def v(message: str, fix: str) -> Violation: return Violation(rule_id, check, message, fix) @@ -386,12 +394,13 @@ def resolve_ci(rule: Rule, root: Path) -> list[Violation]: def resolve_proof_ci(rule: Rule, root: Path) -> list[Violation]: - """A proof test must run in CI, or nothing shows the checker still works. + """Check that a CI job collects the rule's proof test. - The proof test itself asserts that the bad fixture is rejected and the good - one accepted; the guard checks that some CI job collects it. A proof runs - in the rule's own CI job when the rule is pytest-enforced; otherwise the - proof names its job with ``workflow``, ``job`` and ``invocation``. + Without that, nothing shows the checker still works. The proof test itself + asserts that the bad fixture is rejected and the good one accepted; the + guard checks that some CI job collects it. A proof runs in the rule's own + CI job when the rule is pytest-enforced; otherwise the proof names its job + with ``workflow``, ``job`` and ``invocation``. """ if rule.proof is None or not isinstance(rule.proof.fields.get("test"), str): return [] @@ -434,7 +443,11 @@ def _test_coverage( root: Path, check: str, ) -> list[Violation]: - """The job's test command actually collects these root-relative test nodes.""" + """Check that the job's test command collects these root-relative test nodes. + + Only ``make`` targets and ``unittest discover`` commands are inspected; any + other command is trusted. + """ violations = [] collected: list[str] | None = None described = "" diff --git a/scripts/rules/validate.py b/scripts/rules/validate.py index f7d117d65..406b40ef6 100644 --- a/scripts/rules/validate.py +++ b/scripts/rules/validate.py @@ -55,6 +55,7 @@ def baseline_entries(path: Path) -> int: + """Count the entries in a baseline file, skipping blank and ``#`` comment lines.""" return sum( 1 for line in path.read_text().splitlines() @@ -242,6 +243,10 @@ def _validate_proof(rule: Rule, root: Path) -> list[Violation]: def validate(registry: Registry, root: Path, today: date) -> list[Violation]: + """Check the registry's fields and file references against the tree at ``root``. + + ``today`` decides whether review dates and exceptions have expired. + """ violations: list[Violation] = [] def v(rule_id: str | None, message: str, fix: str, check: str = "schema") -> None: diff --git a/scripts/rules/workflows.py b/scripts/rules/workflows.py index a432455fd..5c7c039b4 100644 --- a/scripts/rules/workflows.py +++ b/scripts/rules/workflows.py @@ -32,11 +32,14 @@ class RunBlock: @property def single_line(self) -> bool: + """Whether the block is at most one line of code with no chained commands.""" return self.code_lines <= 1 and not self.chained @dataclass(frozen=True) class JobBlock: + """One workflow job: its run text, step names, ``uses:`` references and matrix.""" + id: str name: str run_text: str diff --git a/scripts/run_quality_guards.py b/scripts/run_quality_guards.py index 5da7c5fb5..a89b7bd3b 100644 --- a/scripts/run_quality_guards.py +++ b/scripts/run_quality_guards.py @@ -24,6 +24,7 @@ def main() -> int: + """Run every guard, print its violations, and return 1 if any guard failed.""" failed = False for name, check in GUARDS: violations = check() diff --git a/scripts/smoke_test_local_boot.py b/scripts/smoke_test_local_boot.py index 5dc8fa2aa..eca2fd272 100644 --- a/scripts/smoke_test_local_boot.py +++ b/scripts/smoke_test_local_boot.py @@ -26,18 +26,21 @@ def find_free_port() -> int: + """Return a localhost TCP port that was free when checked (nothing reserves it).""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) return int(sock.getsockname()[1]) def read_text(path: Path) -> str: + """Return the file's text with undecodable bytes replaced, or "" if missing.""" if not path.exists(): return "" return path.read_text(encoding="utf-8", errors="replace") def fetch(url: str) -> tuple[int, str]: + """Return the status and body of a GET; error statuses raise urllib's HTTPError.""" request = urllib.request.Request(url, method="GET") with urllib.request.urlopen(request, timeout=HTTP_TIMEOUT_SECONDS) as response: body = response.read().decode("utf-8", errors="replace") @@ -45,11 +48,16 @@ def fetch(url: str) -> tuple[int, str]: def print_section(title: str, body: str) -> None: + """Print a titled block of captured output, or "(empty)" when there is none.""" print(f"\n=== {title} ===") print(body.rstrip() if body else "(empty)") def main() -> int: + """Boot the Flask app on a free port and return 1 unless both health checks pass. + + On failure print the app's captured stdout and stderr; always stop the app. + """ port = find_free_port() stdout_file = tempfile.NamedTemporaryFile( mode="w+", diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py index d2c71174f..aa63effc3 100644 --- a/scripts/v1_database_migration.py +++ b/scripts/v1_database_migration.py @@ -28,6 +28,12 @@ class DatabaseState(StrEnum): + """Where a v1 database stands against the Alembic revision scripts. + + UNVERSIONED has no version table and INVALID has an empty one; neither is + ever stamped or upgraded automatically. + """ + UNVERSIONED = "unversioned" INVALID = "invalid" PENDING = "pending" @@ -59,6 +65,10 @@ def classify_database_state( current_heads: set[str], script_heads: set[str] | None = None, ) -> DatabaseState: + """Classify a database by its version table and heads. + + Without ``script_heads`` a versioned database is reported as pending, never head. + """ if not version_table_exists: return DatabaseState.UNVERSIONED if not current_heads: @@ -102,6 +112,7 @@ def _script_heads(config: Config) -> set[str]: def database_state(connection: Connection) -> DatabaseState: + """Return the database's Alembic state against the repository's revision heads.""" inspector = inspect(connection) version_table_exists = "alembic_version" in inspector.get_table_names() context = MigrationContext.configure(connection) @@ -114,6 +125,10 @@ def database_state(connection: Connection) -> DatabaseState: def metadata_differences(connection: Connection) -> set[str]: + """Describe every difference between the live schema and the v1 models. + + Raise ValueError for a kind of difference that has no data-free description. + """ context = MigrationContext.configure( connection, opts={"compare_type": True, "compare_server_default": True}, @@ -125,6 +140,7 @@ def metadata_differences(connection: Connection) -> set[str]: def verify_head_schema(connection: Connection) -> None: + """Raise RuntimeError unless the database is at the head with no schema drift.""" state = database_state(connection) if state is not DatabaseState.HEAD: raise RuntimeError(f"Database is not at the Alembic head; found {state}") @@ -150,6 +166,11 @@ def _release_lock(connection: Connection) -> None: def upgrade_database(connection: Connection, *, backup_id: str) -> None: + """Upgrade a versioned database to head under a MySQL named lock, then verify it. + + A database already at head is only verified, so ``backup_id`` is required + only when an upgrade will actually run. + """ _acquire_lock(connection) try: state = database_state(connection) @@ -209,6 +230,11 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: def main(argv: list[str] | None = None) -> int: + """Run one --mode against the v1 database: state, verify-head or upgrade. + + Unless a database URL is set, the read-only modes connect through the proxy + as the schema reader and upgrade as the migrator, in one transaction. + """ args = _parse_args(argv) engine = create_engine(_database_target(args.mode), poolclass=NullPool) try: From f64e90b6852e65d6d746b1c1f4703738b38572b9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:23:44 +0400 Subject: [PATCH 05/12] Document the data layer Add docstrings to the v1 SQLAlchemy models and connection settings, the congressional-district metadata, the v2 catalog publication records and staging, the household and policy migration qualification, the v2 migration target settings, and the v2 SQLModel tables and their enums. The table docstrings say what a row represents: what makes it unique, whether it is immutable, what owns it and what is deleted with it. Pydantic publishes a class docstring as the model's JSON-schema description, so these tables' schemas gain descriptions; none of them is in the published OpenAPI document. Co-Authored-By: Claude Opus 5.5 --- .../data/congressional_districts.py | 14 +++-- policyengine_api/data/orm.py | 5 ++ policyengine_api/data/v1_models.py | 53 +++++++++++++++- .../v2/catalog/publication_reconciliation.py | 15 +++++ .../data/v2/catalog/publication_staging.py | 6 ++ .../data/v2/catalog/publication_types.py | 1 + policyengine_api/data/v2/catalog/records.py | 33 ++++++++++ .../v2/household_migration_qualification.py | 7 +++ policyengine_api/data/v2/migration_target.py | 10 +++ .../data/v2/models/associations.py | 22 +++++++ .../data/v2/models/comparison_runs.py | 17 ++++++ policyengine_api/data/v2/models/households.py | 14 +++++ policyengine_api/data/v2/models/metadata.py | 49 +++++++++++++++ policyengine_api/data/v2/models/policies.py | 11 ++++ policyengine_api/data/v2/models/reports.py | 61 +++++++++++++++++++ .../data/v2/models/simulations.py | 10 +++ policyengine_api/data/v2/models/users.py | 6 ++ .../data/v2/policy_migration_qualification.py | 7 +++ policyengine_api/data/v2/settings.py | 1 + 19 files changed, 335 insertions(+), 7 deletions(-) diff --git a/policyengine_api/data/congressional_districts.py b/policyengine_api/data/congressional_districts.py index 38a826cc7..5c6603260 100644 --- a/policyengine_api/data/congressional_districts.py +++ b/policyengine_api/data/congressional_districts.py @@ -1,4 +1,4 @@ -"""US Congressional District Metadata +"""Metadata for every US congressional district and helpers that use it. This module defines the metadata for all 435 US Congressional districts based on the 2020 Census apportionment, effective for the 118th Congress @@ -620,7 +620,7 @@ def _format_district_number(number: int) -> str: def _build_district_name(state_code: str, number: int) -> str: - """Build the district name in the format: congressional_district/- + """Build a district's region name, ``congressional_district/-``. Examples: ("CA", 5) -> "congressional_district/CA-05" @@ -632,8 +632,10 @@ def _build_district_name(state_code: str, number: int) -> str: def _build_district_label(state_code: str, number: int) -> str: - """Build the district label in the format: 's th congressional district - For at-large districts (states with only 1 district), use: 's at-large congressional district + """Build a district's display label, such as "Texas's 3rd congressional district". + + At-large districts (states with only one district) are labelled + "'s at-large congressional district" instead of numbered. Examples: ("CA", 1) -> "California's 1st congressional district" @@ -652,7 +654,7 @@ def _build_district_label(state_code: str, number: int) -> str: def build_congressional_district_metadata() -> list[dict]: - """Build the complete congressional district metadata structure for use in country.py. + """Build the congressional district region entries that country.py publishes. Returns a list of dictionaries with 'name', 'label', 'type', 'state_abbreviation', and 'state_name' keys, formatted as: @@ -701,7 +703,7 @@ def get_valid_state_codes() -> set[str]: def get_valid_congressional_districts() -> set[str]: - """Get the set of valid congressional district identifiers (lowercase for case-insensitive matching). + """Return every valid district identifier, lowercased for case-insensitive matching. Format: "-" (e.g., "ca-37", "tx-01") diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 57f345e2d..12e9a5ec2 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -29,6 +29,11 @@ def get_remote_database_config() -> dict[str, str]: + """Read the Cloud SQL instance, user and database name from the environment. + + Raise KeyError when POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is unset; the + user and database name default to ``policyengine``. + """ return { "instance_connection_name": os.environ[ "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME" diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py index 79f83a173..71578d8aa 100644 --- a/policyengine_api/data/v1_models.py +++ b/policyengine_api/data/v1_models.py @@ -28,10 +28,12 @@ class V1Base(DeclarativeBase): - pass + """Declarative base whose metadata is the v1 schema Alembic compares against.""" class Household(V1Base): + """A saved v1 household input, with its JSON hash and country package version.""" + __tablename__ = "household" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3)) @@ -95,6 +97,11 @@ class HouseholdMirrorEvent(V1Base): class ComputedHousehold(V1Base): + """A legacy cached household calculation, keyed by household, policy and country. + + Mapped so Alembic sees the table; the API no longer reads or writes it. + """ + __tablename__ = "computed_household" household_id: Mapped[int] = mapped_column(Integer, primary_key=True) policy_id: Mapped[int] = mapped_column(Integer, primary_key=True) @@ -105,6 +112,8 @@ class ComputedHousehold(V1Base): class Policy(V1Base): + """A saved v1 policy reform, identified within a country by its id and JSON hash.""" + __tablename__ = "policy" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3), primary_key=True) @@ -115,6 +124,11 @@ class Policy(V1Base): class Economy(V1Base): + """A legacy cached economy-wide result for one policy and options hash. + + Mapped so Alembic sees the table; the API no longer reads or writes it. + """ + __tablename__ = "economy" economy_id: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True @@ -132,6 +146,12 @@ class Economy(V1Base): class ReformImpact(V1Base): + """A cached economy comparison of a reform against a baseline, with its job status. + + Rows are keyed by policies, region, dataset, time period, options hash and + API version, and record the simulation execution that computes them. + """ + __tablename__ = "reform_impact" reform_impact_id: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True @@ -154,6 +174,11 @@ class ReformImpact(V1Base): class Analysis(V1Base): + """A legacy AI-analysis prompt and its generated text. + + Mapped so Alembic sees the table; the API no longer reads or writes it. + """ + __tablename__ = "analysis" prompt_id: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True @@ -166,6 +191,12 @@ class Analysis(V1Base): class UserPolicy(V1Base): + """A user's saved reform/baseline pair and how they set it up. + + ``mirror_revision`` increases on every change so v2 mirroring applies + events in order. + """ + __tablename__ = "user_policies" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3)) @@ -253,6 +284,8 @@ class UserPolicyMirrorEvent(V1Base): class UserProfile(V1Base): + """A user's profile, unique by Auth0 id and by username when one is set.""" + __tablename__ = "user_profiles" user_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) auth0_id: Mapped[str] = mapped_column(String(255), unique=True) @@ -262,6 +295,12 @@ class UserProfile(V1Base): class Simulation(V1Base): + """A v1 simulation of one policy on a household or geography population. + + Holds the latest output inline, plus pointers to its active and latest + successful runs in ``simulation_runs``. + """ + __tablename__ = "simulations" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3)) @@ -279,6 +318,12 @@ class Simulation(V1Base): class ReportOutput(V1Base): + """A v1 report comparing one or two simulations, with its latest output inline. + + Also stores the report spec and pointers to its active and latest + successful runs in ``report_output_runs``. + """ + __tablename__ = "report_outputs" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3)) @@ -298,6 +343,8 @@ class ReportOutput(V1Base): class ReportOutputRun(V1Base): + """One numbered execution of a report output, with the versions it ran against.""" + __tablename__ = "report_output_runs" __table_args__ = ( UniqueConstraint( @@ -332,6 +379,8 @@ class ReportOutputRun(V1Base): class SimulationRun(V1Base): + """One numbered execution of a simulation, possibly as an input to a report run.""" + __tablename__ = "simulation_runs" __table_args__ = ( UniqueConstraint( @@ -366,6 +415,8 @@ class SimulationRun(V1Base): class LegacyReportOutputAlias(V1Base): + """Redirect from a duplicate legacy report output id to its canonical one.""" + __tablename__ = "legacy_report_output_aliases" legacy_report_output_id: Mapped[int] = mapped_column( Integer, diff --git a/policyengine_api/data/v2/catalog/publication_reconciliation.py b/policyengine_api/data/v2/catalog/publication_reconciliation.py index 75d6ede2a..3296bf172 100644 --- a/policyengine_api/data/v2/catalog/publication_reconciliation.py +++ b/policyengine_api/data/v2/catalog/publication_reconciliation.py @@ -48,6 +48,7 @@ def version_exists(connection: Connection, country: CountryCatalog) -> bool: + """Return whether the model name and version of this country are already stored.""" statement = ( sa.select(MODEL_VERSIONS.c.id) .select_from( @@ -284,6 +285,11 @@ def assert_country_matches( connection: Connection, country: CountryCatalog, ) -> None: + """Raise CatalogPublicationError if the stored version differs from the staged one. + + Rows are compared as sets in both directions, so a missing row and an extra + row both fail. + """ for staged, actual in _comparison_pairs(country): if _sets_differ(connection, staged, actual): raise CatalogPublicationError( @@ -626,11 +632,20 @@ def assert_country_matches( def publish_new_country(connection: Connection, country: CountryCatalog) -> None: + """Insert one country's staged catalog into the v2 tables, parents first. + + The model row is reused if one with the same name already exists. + """ for statement in SET_BASED_INSERT_STATEMENTS: connection.execute(statement, {"country_id": country.country_id}) def assert_canonical_value_uniqueness(connection: Connection) -> None: + """Raise CatalogPublicationError if a parameter has duplicate current-law values. + + Current-law values are those with neither a policy nor a dynamic; each + parameter may have only one per start date. + """ duplicates = ( sa.select(PARAMETER_VALUES.c.parameter_id) .where( diff --git a/policyengine_api/data/v2/catalog/publication_staging.py b/policyengine_api/data/v2/catalog/publication_staging.py index 2d2f556b0..8656b73f8 100644 --- a/policyengine_api/data/v2/catalog/publication_staging.py +++ b/policyengine_api/data/v2/catalog/publication_staging.py @@ -290,6 +290,7 @@ def copy_rows( def create_staging_tables(connection: Connection) -> None: + """Create the temporary staging tables, which Postgres drops at commit.""" STAGING_METADATA.create_all(connection, checkfirst=False) @@ -299,6 +300,11 @@ def stage_catalog( *, checkpoint: Callable[[str, Connection], None] | None = None, ) -> dict[str, int]: + """COPY every country's records into the staging tables and return rows per table. + + Raise CatalogPublicationError if the copied counts differ from the + catalog's own counts. + """ observed = {table.name: 0 for table in STAGING_TABLES} for country in catalog.countries: for table, rows in _catalog_rows(country).items(): diff --git a/policyengine_api/data/v2/catalog/publication_types.py b/policyengine_api/data/v2/catalog/publication_types.py index b9e388fc0..884bcef07 100644 --- a/policyengine_api/data/v2/catalog/publication_types.py +++ b/policyengine_api/data/v2/catalog/publication_types.py @@ -20,6 +20,7 @@ class PublicationEvidence: elapsed_seconds: float def as_dict(self) -> dict[str, object]: + """Return the evidence as the JSON object printed after a publication.""" return { "outcome": "ok", "policyengine_version": self.policyengine_version, diff --git a/policyengine_api/data/v2/catalog/records.py b/policyengine_api/data/v2/catalog/records.py index b8a5864fe..555311294 100644 --- a/policyengine_api/data/v2/catalog/records.py +++ b/policyengine_api/data/v2/catalog/records.py @@ -13,6 +13,8 @@ @dataclass(frozen=True, slots=True) class ModelRecord: + """One country's tax-benefit model as named in the catalog.""" + id: UUID country_id: str name: str @@ -21,6 +23,12 @@ class ModelRecord: @dataclass(frozen=True, slots=True) class ModelVersionRecord: + """The model version a catalog was extracted from. + + Also carries the country's current-law policy id and the years its metadata + covers. + """ + id: UUID model_id: UUID version: str @@ -31,6 +39,8 @@ class ModelVersionRecord: @dataclass(frozen=True, slots=True) class VariableRecord: + """One model variable, with its entity, type and the variables it sums.""" + id: UUID model_version_id: UUID name: str @@ -46,6 +56,8 @@ class VariableRecord: @dataclass(frozen=True, slots=True) class ParameterNodeRecord: + """A non-leaf node of the parameter tree; it groups parameters and has no values.""" + id: UUID model_version_id: UUID name: str @@ -55,6 +67,8 @@ class ParameterNodeRecord: @dataclass(frozen=True, slots=True) class ParameterValueRecord: + """A parameter's value from ``start_date`` until ``end_date``, or indefinitely.""" + id: UUID parameter_id: UUID value_json: Any @@ -64,6 +78,8 @@ class ParameterValueRecord: @dataclass(frozen=True, slots=True) class ParameterRecord: + """A leaf parameter with its unit and every value it has held over time.""" + id: UUID model_version_id: UUID name: str @@ -76,6 +92,8 @@ class ParameterRecord: @dataclass(frozen=True, slots=True) class DatasetRecord: + """A logical input dataset, referenced by name; ``storage_path`` is always None.""" + id: UUID model_version_id: UUID name: str @@ -87,6 +105,12 @@ class DatasetRecord: @dataclass(frozen=True, slots=True) class RegionRecord: + """A region that can be simulated, with the dataset it uses by default. + + Regions that need a subset of that dataset (``requires_filter``) also name + the variable, value and strategy that select their households. + """ + id: UUID model_version_id: UUID default_dataset_id: UUID @@ -104,6 +128,11 @@ class RegionRecord: @dataclass(frozen=True, slots=True) class FallbackSummary: + """How many US sub-national regions of one type fell back to the default dataset. + + These regions had no dataset path of their own. + """ + region_type: str count: int @@ -125,6 +154,8 @@ def iter_batches( @dataclass(frozen=True, slots=True) class CountryCatalog: + """Every record the catalog publishes for one country's model version.""" + country_id: str model: ModelRecord model_version: ModelVersionRecord @@ -171,6 +202,8 @@ def entity_counts(self) -> dict[str, int]: @dataclass(frozen=True, slots=True) class NormalizedCatalog: + """A complete extracted catalog for one PolicyEngine.py version, all countries.""" + policyengine_version: str dependency_versions: tuple[tuple[str, str], ...] countries: tuple[CountryCatalog, ...] diff --git a/policyengine_api/data/v2/household_migration_qualification.py b/policyengine_api/data/v2/household_migration_qualification.py index c680c8eca..0297a7d85 100644 --- a/policyengine_api/data/v2/household_migration_qualification.py +++ b/policyengine_api/data/v2/household_migration_qualification.py @@ -54,9 +54,11 @@ class HouseholdDataCounts: @property def total(self) -> int: + """The number of affected rows; anything but zero stops the migration.""" return sum(self.as_dict().values()) def as_dict(self) -> dict[str, int]: + """Return the counts keyed by the names used in evidence and errors.""" return { "households": self.households, "user_household_associations": self.user_household_associations, @@ -76,6 +78,10 @@ class HouseholdMigrationQualification: required: bool = True def as_dict(self) -> dict[str, object]: + """Return the evidence as the JSON object the qualification command prints. + + ``qualification`` is "not-required" when the revision was already applied. + """ return { "outcome": "ok", "qualification": "performed" if self.required else "not-required", @@ -89,6 +95,7 @@ class RetainedHouseholdDataError(RuntimeError): """Raised when a target contains household data requiring preservation.""" def __init__(self, counts: HouseholdDataCounts) -> None: + """Keep the counts and say in the message how to proceed safely.""" self.counts = counts rendered_counts = ", ".join( f"{name}={count}" for name, count in counts.as_dict().items() diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index bc3b9d646..1fbc8fd71 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -36,6 +36,11 @@ class V2MigrationTargetError(V2ConfigurationError): @dataclass(frozen=True) class ConfiguredSupabaseTarget: + """The persistent Supabase database v2 Alembic may migrate, and its required role. + + The freshness audit fields gate the first migration of an unstamped database. + """ + environment: str project_ref: str database_name: str @@ -46,6 +51,11 @@ class ConfiguredSupabaseTarget: @dataclass(frozen=True) class V2AlembicSettings: + """The URL v2 Alembic connects to, and whether it is the disposable test database. + + ``target`` is None only in disposable test mode. + """ + url: URL disposable_test: bool target: ConfiguredSupabaseTarget | None diff --git a/policyengine_api/data/v2/models/associations.py b/policyengine_api/data/v2/models/associations.py index 1f0c32de8..df2f39a18 100644 --- a/policyengine_api/data/v2/models/associations.py +++ b/policyengine_api/data/v2/models/associations.py @@ -21,6 +21,11 @@ class UserHouseholdAssociation(TimestampedModel, table=True): + """A user's saved, named link to a household in the same country. + + Deleted with the user; a household that users have saved cannot be deleted. + """ + __tablename__ = "user_household_associations" __table_args__ = ( sa.UniqueConstraint( @@ -67,6 +72,12 @@ class UserHouseholdAssociation(TimestampedModel, table=True): class UserPolicy(TimestampedModel, table=True): + """A user's saved, named link to a policy in the same country. + + Deleted with the user, together with any v1 mapping it was mirrored from; a + policy that users have saved cannot be deleted. + """ + __tablename__ = "user_policies" __table_args__ = ( sa.UniqueConstraint( @@ -117,6 +128,11 @@ class UserPolicy(TimestampedModel, table=True): class UserSimulationAssociation(TimestampedModel, table=True): + """A user's labelled link to a simulation, at most one per user and simulation. + + Deleted with either the user or the simulation. + """ + __tablename__ = "user_simulation_associations" __table_args__ = ( sa.UniqueConstraint( @@ -144,6 +160,12 @@ class UserSimulationAssociation(TimestampedModel, table=True): class UserReportAssociation(TimestampedModel, table=True): + """A user's labelled link to a report, at most one per user and report. + + Records when the user last ran the report, and is deleted with either the + user or the report. + """ + __tablename__ = "user_report_associations" __table_args__ = ( sa.UniqueConstraint( diff --git a/policyengine_api/data/v2/models/comparison_runs.py b/policyengine_api/data/v2/models/comparison_runs.py index da24b6886..e45e992af 100644 --- a/policyengine_api/data/v2/models/comparison_runs.py +++ b/policyengine_api/data/v2/models/comparison_runs.py @@ -11,6 +11,11 @@ class Stage12RunStatus(StrEnum): + """Lifecycle state of a Stage 12 run and of each of its simulations. + + Failed and skipped rows always carry an error code. + """ + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -20,6 +25,11 @@ class Stage12RunStatus(StrEnum): class Stage12AggregationStatus(StrEnum): + """Progress of combining a Stage 12 run's simulations into its aggregate. + + A run can only succeed once its aggregation has succeeded. + """ + NOT_STARTED = "not_started" RUNNING = "running" SUCCEEDED = "succeeded" @@ -27,6 +37,11 @@ class Stage12AggregationStatus(StrEnum): class Stage12ResultComparisonStatus(StrEnum): + """Outcome of comparing a Stage 12 aggregate exactly with the production result. + + Direct submissions have no production result, so they stay ``not_requested``. + """ + NOT_REQUESTED = "not_requested" PENDING = "pending" RUNNING = "running" @@ -36,6 +51,8 @@ class Stage12ResultComparisonStatus(StrEnum): class Stage12SimulationRole(StrEnum): + """Role of a Stage 12 simulation in its run: baseline, reform or standalone.""" + BASELINE = "baseline" REFORM = "reform" STANDALONE = "standalone" diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index a6831d659..78e76b2ab 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -25,6 +25,8 @@ class HouseholdJobStatus(StrEnum): + """Lifecycle state of a queued household calculation job.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -32,6 +34,13 @@ class HouseholdJobStatus(StrEnum): class Household(TimestampedModel, table=True): + """Immutable household input, stored once per canonicalization version and hash. + + ``content_hash`` is the SHA-256 of the canonical document, so saving the same + household again finds the existing row. Households that simulations, reports + or saved links use cannot be deleted. + """ + __tablename__ = "households" __table_args__ = ( sa.UniqueConstraint( @@ -88,6 +97,11 @@ class Household(TimestampedModel, table=True): class HouseholdJob(IdentifiedModel, table=True): + """A queued household calculation request, with its status and result or error. + + Deleting the job's policy or dynamic keeps the job but clears the reference. + """ + __tablename__ = "household_jobs" __table_args__ = ( sa.Index("ix_household_jobs_status_created_at", "status", "created_at"), diff --git a/policyengine_api/data/v2/models/metadata.py b/policyengine_api/data/v2/models/metadata.py index 1b7b9edb7..8d91864df 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -21,6 +21,11 @@ class RegionType(StrEnum): + """Kind of geography a region covers. + + ``national`` is the whole model country; ``country`` is a UK nation. + """ + NATIONAL = "national" COUNTRY = "country" STATE = "state" @@ -32,6 +37,11 @@ class RegionType(StrEnum): class TaxBenefitModel(TimestampedModel, table=True): + """A country's tax-benefit model, unique by name. + + Its versions own the metadata catalog and are deleted with it. + """ + __tablename__ = "tax_benefit_models" __table_args__ = (sa.UniqueConstraint("name", name="uq_tax_benefit_models_name"),) @@ -50,6 +60,13 @@ class TaxBenefitModel(TimestampedModel, table=True): class TaxBenefitModelVersion(IdentifiedModel, table=True): + """One published catalog version of a model, unique per model and version. + + Owns its variables, parameters and parameter nodes, which are deleted with it, + and records the country's current-law policy id and the years its metadata + covers. + """ + __tablename__ = "tax_benefit_model_versions" __table_args__ = ( sa.UniqueConstraint( @@ -91,6 +108,13 @@ class TaxBenefitModelVersion(IdentifiedModel, table=True): class Region(TimestampedModel, table=True): + """A region in one model version, unique by code within that version. + + Its default dataset must belong to the same version. Filtered regions + (``requires_filter``) also name the variable and value that select their + households. + """ + __tablename__ = "regions" __table_args__ = ( sa.UniqueConstraint( @@ -142,6 +166,11 @@ class Region(TimestampedModel, table=True): class Dataset(TimestampedModel, table=True): + """A dataset of one model version, unique by name within that version. + + Output datasets (``is_output_dataset``) must record a ``storage_path``. + """ + __tablename__ = "datasets" __table_args__ = ( sa.UniqueConstraint( @@ -198,6 +227,8 @@ class Dataset(TimestampedModel, table=True): class DatasetVersion(IdentifiedModel, table=True): + """A named version of a dataset, unique per dataset and deleted with it.""" + __tablename__ = "dataset_versions" __table_args__ = ( sa.UniqueConstraint( @@ -225,6 +256,8 @@ class DatasetVersion(IdentifiedModel, table=True): class Variable(IdentifiedModel, table=True): + """A model variable in one model version, unique by name and deleted with it.""" + __tablename__ = "variables" __table_args__ = ( sa.UniqueConstraint( @@ -255,6 +288,11 @@ class Variable(IdentifiedModel, table=True): class ParameterNode(IdentifiedModel, table=True): + """A non-leaf node of a version's parameter tree, which groups parameters. + + Unique by name within the version, and has no values of its own. + """ + __tablename__ = "parameter_nodes" __table_args__ = ( sa.UniqueConstraint( @@ -279,6 +317,11 @@ class ParameterNode(IdentifiedModel, table=True): class Parameter(IdentifiedModel, table=True): + """A leaf parameter of a model version, unique by name within it. + + Its values include current law and the values each policy or dynamic sets. + """ + __tablename__ = "parameters" __table_args__ = ( sa.UniqueConstraint( @@ -309,6 +352,12 @@ class Parameter(IdentifiedModel, table=True): class ParameterValue(IdentifiedModel, table=True): + """A parameter's value from ``start_date`` until ``end_date``, or indefinitely. + + A value owned by a policy or dynamic is deleted with its owner. A value with + neither is current law, and a parameter has at most one per start date. + """ + __tablename__ = "parameter_values" __table_args__ = ( sa.CheckConstraint( diff --git a/policyengine_api/data/v2/models/policies.py b/policyengine_api/data/v2/models/policies.py index 9b9957b30..97e4a619e 100644 --- a/policyengine_api/data/v2/models/policies.py +++ b/policyengine_api/data/v2/models/policies.py @@ -22,6 +22,12 @@ class Policy(TimestampedModel, table=True): + """Immutable policy content, stored once per canonicalization version and hash. + + Its parameter values are deleted with it. Policies that saved links use + cannot be deleted. + """ + __tablename__ = "policies" __table_args__ = ( sa.UniqueConstraint( @@ -90,6 +96,11 @@ class Policy(TimestampedModel, table=True): class Dynamic(TimestampedModel, table=True): + """Named parameter values, separate from any policy, that a simulation can apply. + + Its parameter values are deleted with it. + """ + __tablename__ = "dynamics" name: str = Field(max_length=255) diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 1213ef907..29b8f0aa5 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -26,6 +26,8 @@ class ReportRunStatus(StrEnum): + """Lifecycle state of a report run; finished runs record a completion time.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -33,12 +35,16 @@ class ReportRunStatus(StrEnum): class ReportRunTrigger(StrEnum): + """Why a report run started: the report's first run, a manual or a system rerun.""" + INITIAL = "initial" MANUAL = "manual" SYSTEM = "system" class OutputStatus(StrEnum): + """Lifecycle state of one computed output row of a report run.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -46,12 +52,16 @@ class OutputStatus(StrEnum): class AggregateType(StrEnum): + """How an aggregate combines a variable's values.""" + SUM = "sum" MEAN = "mean" COUNT = "count" class DecileType(StrEnum): + """Whether deciles rank by income or by wealth.""" + INCOME = "income" WEALTH = "wealth" @@ -237,6 +247,11 @@ class ReportRun(TimestampedModel, table=True): class AggregateOutput(IdentifiedModel, table=True): + """A sum, mean or count of one variable in one simulation of a report run. + + ``filter_config`` optionally restricts which entities are included. + """ + __tablename__ = "aggregates" __table_args__ = ( sa.Index("ix_aggregates_report_run_status", "report_run_id", "status"), @@ -267,6 +282,12 @@ class AggregateOutput(IdentifiedModel, table=True): class ChangeAggregate(IdentifiedModel, table=True): + """An aggregate of a variable's change from baseline to reform in a report run. + + ``change_geq`` and ``change_leq`` optionally restrict it to entities whose + change falls within those bounds. + """ + __tablename__ = "change_aggregates" __table_args__ = ( sa.Index( @@ -307,6 +328,11 @@ class ChangeAggregate(IdentifiedModel, table=True): class BudgetSummary(IdentifiedModel, table=True): + """Baseline and reform totals of one budget variable for a report run. + + One row per run, variable and entity. + """ + __tablename__ = "budget_summary" __table_args__ = ( sa.UniqueConstraint( @@ -339,6 +365,12 @@ class BudgetSummary(IdentifiedModel, table=True): class DecileImpact(IdentifiedModel, table=True): + """Mean income change and better/worse-off counts in one income quantile. + + ``decile`` runs from 1 to ``quantiles``; there is one row per run, income + variable, entity, quantile count and decile. + """ + __tablename__ = "decile_impacts" __table_args__ = ( sa.UniqueConstraint( @@ -384,6 +416,11 @@ class DecileImpact(IdentifiedModel, table=True): class IntraDecileImpact(IdentifiedModel, table=True): + """Shares of one decile whose income rises or falls by more or less than 5%. + + One row per run, decile type and decile, where ``decile`` is 0 to 10. + """ + __tablename__ = "intra_decile_impacts" __table_args__ = ( sa.UniqueConstraint( @@ -425,6 +462,8 @@ class IntraDecileImpact(IdentifiedModel, table=True): class Inequality(IdentifiedModel, table=True): + """Gini and top/bottom income shares for one simulation of a report run.""" + __tablename__ = "inequality" __table_args__ = ( sa.UniqueConstraint( @@ -455,6 +494,8 @@ class Inequality(IdentifiedModel, table=True): class Poverty(IdentifiedModel, table=True): + """Poverty headcount and rate of one type for one simulation of a report run.""" + __tablename__ = "poverty" __table_args__ = ( sa.Index( @@ -484,6 +525,11 @@ class Poverty(IdentifiedModel, table=True): class ProgramStatistics(IdentifiedModel, table=True): + """Baseline and reform totals, counts, winners and losers for one program. + + One row per run, program and entity; ``is_tax`` marks taxes. + """ + __tablename__ = "program_statistics" __table_args__ = ( sa.UniqueConstraint( @@ -541,6 +587,11 @@ class GeographicImpactBase(IdentifiedModel): class CongressionalDistrictImpact(GeographicImpactBase, table=True): + """Household income change and population in one US congressional district. + + One row per run and district GEOID. + """ + __tablename__ = "congressional_district_impacts" __table_args__ = ( sa.UniqueConstraint( @@ -560,6 +611,11 @@ class CongressionalDistrictImpact(GeographicImpactBase, table=True): class ConstituencyImpact(GeographicImpactBase, table=True): + """Household income change in one UK constituency, with its map coordinates. + + One row per run and constituency code. + """ + __tablename__ = "constituency_impacts" __table_args__ = ( sa.UniqueConstraint( @@ -578,6 +634,11 @@ class ConstituencyImpact(GeographicImpactBase, table=True): class LocalAuthorityImpact(GeographicImpactBase, table=True): + """Household income change in one UK local authority, with its map coordinates. + + One row per run and local authority code. + """ + __tablename__ = "local_authority_impacts" __table_args__ = ( sa.UniqueConstraint( diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index 816d4896f..bd81f3288 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -23,6 +23,8 @@ class SimulationStatus(StrEnum): + """Lifecycle state of a v2 simulation.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -30,11 +32,19 @@ class SimulationStatus(StrEnum): class SimulationType(StrEnum): + """Whether a simulation runs on one household or on a population dataset.""" + HOUSEHOLD = "household" ECONOMY = "economy" class Simulation(TimestampedModel, table=True): + """A simulation of exactly one household or one dataset, matching its type. + + Deleting its policy, dynamic, region or output dataset clears that reference; + its household and input dataset cannot be deleted while it exists. + """ + __tablename__ = "simulations" __table_args__ = ( sa.CheckConstraint( diff --git a/policyengine_api/data/v2/models/users.py b/policyengine_api/data/v2/models/users.py index 87db70d6b..1f05490b3 100644 --- a/policyengine_api/data/v2/models/users.py +++ b/policyengine_api/data/v2/models/users.py @@ -18,6 +18,12 @@ class User(IdentifiedModel, table=True): + """A v2 user account, unique by email when one is given. + + Deleting a user deletes their saved associations; reports they own are kept + without an owner. + """ + __tablename__ = "users" __table_args__ = ( sa.UniqueConstraint("email", name="uq_users_email"), diff --git a/policyengine_api/data/v2/policy_migration_qualification.py b/policyengine_api/data/v2/policy_migration_qualification.py index 5a514549e..597539aff 100644 --- a/policyengine_api/data/v2/policy_migration_qualification.py +++ b/policyengine_api/data/v2/policy_migration_qualification.py @@ -46,9 +46,11 @@ class PolicyDataCounts: @property def total(self) -> int: + """The number of affected rows; anything but zero stops the migration.""" return self.policies + self.policy_parameter_values + self.user_policies def as_dict(self) -> dict[str, int]: + """Return the counts keyed by the names used in evidence and errors.""" return { "policies": self.policies, "policy_parameter_values": self.policy_parameter_values, @@ -66,6 +68,10 @@ class PolicyMigrationQualification: required: bool = True def as_dict(self) -> dict[str, object]: + """Return the evidence as the JSON object the qualification command prints. + + ``qualification`` is "not-required" when the revision was already applied. + """ return { "outcome": "ok", "qualification": "performed" if self.required else "not-required", @@ -79,6 +85,7 @@ class RetainedPolicyDataError(RuntimeError): """Raised when a target contains policy data requiring preservation.""" def __init__(self, counts: PolicyDataCounts) -> None: + """Keep the counts and say in the message how to proceed safely.""" self.counts = counts super().__init__( "the configured Supabase target contains retained v2 policy data " diff --git a/policyengine_api/data/v2/settings.py b/policyengine_api/data/v2/settings.py index 98c852dea..5aa3833db 100644 --- a/policyengine_api/data/v2/settings.py +++ b/policyengine_api/data/v2/settings.py @@ -57,6 +57,7 @@ def redacted_url(self) -> str: return self._url.render_as_string(hide_password=True) def __str__(self) -> str: + """Return the URL with its password hidden, so printing never leaks it.""" return self.redacted_url From 15a01a790e8065a24107d0989b9baded8cf5dd38 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:23:44 +0400 Subject: [PATCH 06/12] Document the v2 services Add docstrings to the v2 resource services (households, policies, user-households and user-policies), the metadata catalog service, and the comparison-run, simulation and report contracts: their sessions, database connectors, transformations, validators and types. They say what each read returns when nothing matches, which writes are deduplicated or idempotent, which typed errors are raised, and how the limit + 1 over-fetch drives has_more. Pydantic publishes class docstrings as JSON-schema descriptions. The metadata result types appear in the published OpenAPI document, so their docstrings are written for API consumers; the simulation and report types are the versioned Stage 12 worker contracts, so theirs describe what each message means to producer and consumer. Co-Authored-By: Claude Opus 5.5 --- .../database_connectors/creates.py | 10 +++ .../database_connectors/reads.py | 18 ++++++ .../database_connectors/updates.py | 12 ++++ .../v2/comparison_runs/database_session.py | 5 ++ .../services/v2/comparison_runs/services.py | 38 ++++++++++++ .../v2/comparison_runs/transformations.py | 13 ++++ .../services/v2/comparison_runs/types.py | 30 +++++++++ .../services/v2/comparison_runs/validators.py | 30 +++++++++ .../households/database_connectors/creates.py | 10 +++ .../households/database_connectors/reads.py | 11 ++++ .../v2/households/database_session.py | 5 ++ .../services/v2/households/services.py | 24 +++++++ .../services/v2/households/transformations.py | 20 ++++++ .../services/v2/households/types.py | 22 +++++++ .../v2/metadata/database_connectors/reads.py | 9 +++ .../database_connectors/reads_datasets.py | 5 ++ .../reads_parameter_tree.py | 7 +++ .../database_connectors/reads_parameters.py | 15 +++++ .../database_connectors/reads_regions.py | 10 +++ .../database_connectors/reads_variables.py | 7 +++ .../services/v2/metadata/database_session.py | 4 ++ .../services/v2/metadata/services.py | 55 +++++++++++++++- .../services/v2/metadata/transformations.py | 23 +++++++ .../services/v2/metadata/types.py | 62 +++++++++++++++++++ .../services/v2/metadata/validators.py | 5 ++ .../policies/database_connectors/creates.py | 12 ++++ .../v2/policies/database_connectors/reads.py | 20 ++++++ .../services/v2/policies/database_session.py | 3 + .../services/v2/policies/services.py | 29 +++++++++ .../services/v2/policies/transformations.py | 25 ++++++++ .../services/v2/policies/types.py | 9 +++ policyengine_api/services/v2/reports/types.py | 22 +++++++ .../services/v2/simulations/types.py | 50 +++++++++++++++ .../database_connectors/creates.py | 1 + .../database_connectors/deletes.py | 1 + .../database_connectors/reads.py | 8 +++ .../database_connectors/updates.py | 1 + .../v2/user_households/database_session.py | 5 ++ .../services/v2/user_households/services.py | 16 +++++ .../v2/user_households/transformations.py | 5 ++ .../services/v2/user_households/types.py | 15 +++++ .../services/v2/user_households/validators.py | 6 ++ .../database_connectors/creates.py | 12 ++++ .../database_connectors/deletes.py | 2 + .../database_connectors/reads.py | 15 +++++ .../database_connectors/updates.py | 5 ++ .../v2/user_policies/database_session.py | 3 + .../services/v2/user_policies/services.py | 35 +++++++++++ .../v2/user_policies/transformations.py | 10 +++ .../services/v2/user_policies/types.py | 24 +++++++ .../services/v2/user_policies/validators.py | 16 +++++ 51 files changed, 799 insertions(+), 1 deletion(-) diff --git a/policyengine_api/services/v2/comparison_runs/database_connectors/creates.py b/policyengine_api/services/v2/comparison_runs/database_connectors/creates.py index 32c0ca17c..14a0412d1 100644 --- a/policyengine_api/services/v2/comparison_runs/database_connectors/creates.py +++ b/policyengine_api/services/v2/comparison_runs/database_connectors/creates.py @@ -29,6 +29,11 @@ def create_comparison_report( session: Session, record: ComparisonReportRecord, ) -> UUID | None: + """Insert a report row unless a row with the same identity already exists. + + Return the new evaluation_id, or None when uq_stage12_eval_reports_identity + already holds a row for this identity. + """ statement = ( insert(Stage12ComparisonReport) .values(**report_insert_values(record)) @@ -42,6 +47,11 @@ def create_comparison_simulation( session: Session, record: ComparisonSimulationRecord, ) -> UUID | None: + """Insert a simulation row unless a row with the same identity already exists. + + Return the new simulation_execution_id, or None when + uq_stage12_eval_simulations_identity already holds a row for this identity. + """ statement = ( insert(Stage12ComparisonSimulation) .values(**simulation_insert_values(record)) diff --git a/policyengine_api/services/v2/comparison_runs/database_connectors/reads.py b/policyengine_api/services/v2/comparison_runs/database_connectors/reads.py index c06b6d337..032f3feb2 100644 --- a/policyengine_api/services/v2/comparison_runs/database_connectors/reads.py +++ b/policyengine_api/services/v2/comparison_runs/database_connectors/reads.py @@ -26,6 +26,10 @@ def read_comparison_report( *, lock: bool, ) -> Stage12ComparisonReport | None: + """Return the report row for evaluation_id, or None. + + With lock=True the row is read FOR UPDATE until the transaction ends. + """ statement = select(Stage12ComparisonReport).where( Stage12ComparisonReport.evaluation_id == evaluation_id ) @@ -38,6 +42,11 @@ def read_comparison_report_by_identity( session: Session, record: ComparisonReportRecord, ) -> Stage12ComparisonReport | None: + """Return the report row with the record's identity, or None. + + The identity is the environment, calculation flow, production identity, + worker version and version-manifest digest. + """ return session.exec( select(Stage12ComparisonReport).where( Stage12ComparisonReport.environment == record.environment, @@ -56,6 +65,10 @@ def read_comparison_simulation( *, lock: bool, ) -> Stage12ComparisonSimulation | None: + """Return the simulation row for simulation_execution_id, or None. + + With lock=True the row is read FOR UPDATE until the transaction ends. + """ statement = select(Stage12ComparisonSimulation).where( Stage12ComparisonSimulation.simulation_execution_id == simulation_execution_id ) @@ -68,6 +81,11 @@ def read_comparison_simulation_by_identity( session: Session, record: ComparisonSimulationRecord, ) -> Stage12ComparisonSimulation | None: + """Return the simulation row with the record's identity, or None. + + The identity is the parent evaluation, role, input digest, worker version + and version-manifest digest. + """ return session.exec( select(Stage12ComparisonSimulation).where( Stage12ComparisonSimulation.evaluation_id == record.evaluation_id, diff --git a/policyengine_api/services/v2/comparison_runs/database_connectors/updates.py b/policyengine_api/services/v2/comparison_runs/database_connectors/updates.py index 91f4fa773..edfb0b31f 100644 --- a/policyengine_api/services/v2/comparison_runs/database_connectors/updates.py +++ b/policyengine_api/services/v2/comparison_runs/database_connectors/updates.py @@ -60,6 +60,10 @@ def update_comparison_report( row: Stage12ComparisonReport, record: ComparisonReportRecord, ) -> Stage12ComparisonReport: + """Copy the record's statuses and mutable output fields onto the report row. + + The row is flushed and refreshed; committing is left to the caller. + """ row.status = Stage12RunStatus(record.status.value) row.aggregation_status = Stage12AggregationStatus(record.aggregation_status.value) for field_name in REPORT_MUTABLE_FIELDS: @@ -75,6 +79,10 @@ def update_comparison_simulation( row: Stage12ComparisonSimulation, record: ComparisonSimulationRecord, ) -> Stage12ComparisonSimulation: + """Copy the record's status, row identity and mutable output fields onto the row. + + The row is flushed and refreshed; committing is left to the caller. + """ row.status = Stage12RunStatus(record.status.value) row.row_identity_columns = ( list(record.row_identity_columns) @@ -94,6 +102,10 @@ def update_report_result_comparison( row: Stage12ComparisonReport, record: ComparisonReportRecord, ) -> Stage12ComparisonReport: + """Copy the record's result-comparison status and fields onto the report row. + + The row is flushed and refreshed; committing is left to the caller. + """ row.comparison_status = Stage12ResultComparisonStatus( record.comparison_status.value ) diff --git a/policyengine_api/services/v2/comparison_runs/database_session.py b/policyengine_api/services/v2/comparison_runs/database_session.py index e3d156915..04a1ddc53 100644 --- a/policyengine_api/services/v2/comparison_runs/database_session.py +++ b/policyengine_api/services/v2/comparison_runs/database_session.py @@ -13,15 +13,20 @@ class ComparisonRunDatabaseSession: + """Session lifetimes for comparison-run services: reads and transactions.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Open every session from session_factory.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads that is closed, not committed, on exit.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session in a transaction that commits on success, else rolls back.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/comparison_runs/services.py b/policyengine_api/services/v2/comparison_runs/services.py index abb22116a..5d94fad94 100644 --- a/policyengine_api/services/v2/comparison_runs/services.py +++ b/policyengine_api/services/v2/comparison_runs/services.py @@ -61,6 +61,11 @@ def create_or_resolve_report( session: Session, record: ComparisonReportRecord, ) -> ComparisonReportPersistenceResult: + """Insert a report row, or resolve the existing row with the same identity. + + Raise ComparisonRunIdentityError when an existing row disagrees with the + record on its release or data identity, so a retry cannot reinterpret it. + """ created_id = create_comparison_report(session, record) if created_id is not None: row = read_comparison_report(session, created_id, lock=False) @@ -86,6 +91,12 @@ def create_or_resolve_simulation( session: Session, record: ComparisonSimulationRecord, ) -> ComparisonSimulationPersistenceResult: + """Insert a simulation row under its report, or resolve the existing row. + + Raise ComparisonRunNotFoundError when the parent report is missing, and + ComparisonRunIdentityError when the record disagrees with its parent's + release and retention fields or with an existing row of the same identity. + """ parent_row = read_comparison_report(session, record.evaluation_id, lock=False) if parent_row is None: raise ComparisonRunNotFoundError( @@ -117,6 +128,12 @@ def replace_report_lifecycle( session: Session, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write a report's new lifecycle, aggregation and output fields under a row lock. + + Identity and result-comparison fields must be unchanged and both status moves + allowed. A replay against a succeeded report must match it exactly, and then + returns the stored record without writing. + """ row = read_comparison_report(session, record.evaluation_id, lock=True) if row is None: raise ComparisonRunNotFoundError( @@ -140,6 +157,12 @@ def replace_report_result_comparison( session: Session, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write a report's result-comparison outcome under a row lock. + + Only the comparison fields and updated_at may differ from the stored row. Once + the comparison is matched or different, an identical replay returns the + stored record without writing and any change raises. + """ row = read_comparison_report(session, record.evaluation_id, lock=True) if row is None: raise ComparisonRunNotFoundError( @@ -164,6 +187,12 @@ def replace_simulation_lifecycle( session: Session, record: ComparisonSimulationRecord, ) -> ComparisonSimulationRecord: + """Write a simulation's new lifecycle and output fields under a row lock. + + Identity fields must be unchanged and the status move allowed. A replay + against a succeeded simulation must match it exactly, and then returns the + stored record without writing. + """ row = read_comparison_simulation( session, record.simulation_execution_id, @@ -183,13 +212,17 @@ def replace_simulation_lifecycle( class V2ComparisonRunService: + """Comparison-run operations, each run in its own database transaction.""" + def __init__(self, database_session: ComparisonRunDatabaseSession) -> None: + """Open a new session from database_session for every call.""" self._database_session = database_session def create_or_resolve_report( self, record: ComparisonReportRecord, ) -> ComparisonReportPersistenceResult: + """Insert a report row or resolve its existing identity, in one transaction.""" with self._database_session.transaction() as session: return create_or_resolve_report(session, record) @@ -197,6 +230,7 @@ def create_or_resolve_simulation( self, record: ComparisonSimulationRecord, ) -> ComparisonSimulationPersistenceResult: + """Insert a simulation row or resolve its identity, in one transaction.""" with self._database_session.transaction() as session: return create_or_resolve_simulation(session, record) @@ -204,6 +238,7 @@ def replace_report_lifecycle( self, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write a report's lifecycle update in one transaction, under a row lock.""" with self._database_session.transaction() as session: return replace_report_lifecycle(session, record) @@ -211,6 +246,7 @@ def replace_simulation_lifecycle( self, record: ComparisonSimulationRecord, ) -> ComparisonSimulationRecord: + """Write a simulation lifecycle update in one transaction, under a row lock.""" with self._database_session.transaction() as session: return replace_simulation_lifecycle(session, record) @@ -218,10 +254,12 @@ def replace_report_result_comparison( self, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write a report's result-comparison outcome in one transaction.""" with self._database_session.transaction() as session: return replace_report_result_comparison(session, record) def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: + """Return the stored report, or raise ComparisonRunNotFoundError.""" with self._database_session.read() as session: row = read_comparison_report(session, evaluation_id, lock=False) if row is None: diff --git a/policyengine_api/services/v2/comparison_runs/transformations.py b/policyengine_api/services/v2/comparison_runs/transformations.py index 5c0c247d8..fc71b24d3 100644 --- a/policyengine_api/services/v2/comparison_runs/transformations.py +++ b/policyengine_api/services/v2/comparison_runs/transformations.py @@ -19,6 +19,7 @@ def report_insert_values(record: ComparisonReportRecord) -> dict[str, Any]: + """Return a report record's column values, with statuses as the table's enums.""" values = record.model_dump(mode="python") values["status"] = Stage12RunStatus(record.status.value) values["aggregation_status"] = Stage12AggregationStatus( @@ -31,6 +32,10 @@ def report_insert_values(record: ComparisonReportRecord) -> dict[str, Any]: def simulation_insert_values(record: ComparisonSimulationRecord) -> dict[str, Any]: + """Return a simulation record's column values in the table's types. + + Status and role become the table's enums, and row_identity_columns a list. + """ values = record.model_dump(mode="python") values["status"] = Stage12RunStatus(record.status.value) values["role"] = Stage12SimulationRole(record.role.value) @@ -43,6 +48,10 @@ def simulation_insert_values(record: ComparisonSimulationRecord) -> dict[str, An def report_record(row: Stage12ComparisonReport) -> ComparisonReportRecord: + """Rebuild the validated contract record from a stored report row. + + Raise pydantic.ValidationError if the row breaks the contract's state rules. + """ return ComparisonReportRecord.model_validate( { field_name: getattr(row, field_name) @@ -54,6 +63,10 @@ def report_record(row: Stage12ComparisonReport) -> ComparisonReportRecord: def simulation_record( row: Stage12ComparisonSimulation, ) -> ComparisonSimulationRecord: + """Rebuild the validated contract record from a stored simulation row. + + Raise pydantic.ValidationError if the row breaks the contract's state rules. + """ return ComparisonSimulationRecord.model_validate( { field_name: getattr(row, field_name) diff --git a/policyengine_api/services/v2/comparison_runs/types.py b/policyengine_api/services/v2/comparison_runs/types.py index 34be2e22f..00a127455 100644 --- a/policyengine_api/services/v2/comparison_runs/types.py +++ b/policyengine_api/services/v2/comparison_runs/types.py @@ -36,6 +36,11 @@ def _validate_retention(created_at: datetime, retention_expires_at: datetime) -> class ComparisonRunLifecycleStatus(StrEnum): + """Lifecycle of a comparison report or simulation. + + Succeeded and skipped are final; failed and incomplete runs may restart. + """ + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -45,6 +50,8 @@ class ComparisonRunLifecycleStatus(StrEnum): class ComparisonRunAggregationStatus(StrEnum): + """Progress of a comparison report's aggregation step; succeeded is final.""" + NOT_STARTED = "not_started" RUNNING = "running" SUCCEEDED = "succeeded" @@ -52,6 +59,11 @@ class ComparisonRunAggregationStatus(StrEnum): class ResultComparisonStatus(StrEnum): + """Outcome of comparing a comparison run's aggregate with the production result. + + Not-requested, matched, different and failed are final. + """ + NOT_REQUESTED = "not_requested" PENDING = "pending" RUNNING = "running" @@ -106,6 +118,13 @@ class ComparisonReportRecord(StrictContractModel): @model_validator(mode="after") def validate_state(self) -> ComparisonReportRecord: + """Reject a report whose lifecycle, output and comparison fields disagree. + + The aggregate output fields are all set or all unset; a succeeded report + needs successful aggregation, its output and completed_at; a failed or + skipped one needs an error_code; and the comparison fields must fit + comparison_status. Retention must be timezone-aware and 1-30 days long. + """ _validate_retention(self.created_at, self.retention_expires_at) output_fields = ( self.aggregate_output_uri, @@ -208,6 +227,13 @@ class ComparisonSimulationRecord(StrictContractModel): @model_validator(mode="after") def validate_state(self) -> ComparisonSimulationRecord: + """Reject a simulation whose lifecycle and output fields disagree. + + The output artifact and row-identity fields are all set or all unset, + row_identity_columns is non-empty and unique, a succeeded simulation + needs its output and completed_at, and a failed or skipped one needs an + error_code. Retention must be timezone-aware and 1-30 days long. + """ _validate_retention(self.created_at, self.retention_expires_at) output_fields = ( self.output_uri, @@ -245,11 +271,15 @@ def validate_state(self) -> ComparisonSimulationRecord: @dataclass(frozen=True) class ComparisonReportPersistenceResult: + """A stored report record, and whether this call inserted it or found it.""" + record: ComparisonReportRecord created: bool @dataclass(frozen=True) class ComparisonSimulationPersistenceResult: + """A stored simulation record, and whether this call inserted it or found it.""" + record: ComparisonSimulationRecord created: bool diff --git a/policyengine_api/services/v2/comparison_runs/validators.py b/policyengine_api/services/v2/comparison_runs/validators.py index 2e5fafa69..45c104a9d 100644 --- a/policyengine_api/services/v2/comparison_runs/validators.py +++ b/policyengine_api/services/v2/comparison_runs/validators.py @@ -135,6 +135,7 @@ def require_report_identity( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ComparisonRunIdentityError if an update changes an immutable field.""" _require_equal_fields(existing, candidate, REPORT_IMMUTABLE_FIELDS) @@ -142,6 +143,11 @@ def require_report_conflict_matches( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ComparisonRunIdentityError if a duplicate insert names a different release. + + The stored report and the retried record must agree on every release, + package and dataset field. + """ _require_equal_fields(existing, candidate, REPORT_CONFLICT_FIELDS) @@ -149,6 +155,7 @@ def require_simulation_identity( existing: ComparisonSimulationRecord, candidate: ComparisonSimulationRecord, ) -> None: + """Raise ComparisonRunIdentityError if an update changes an immutable field.""" _require_equal_fields(existing, candidate, SIMULATION_IMMUTABLE_FIELDS) @@ -156,6 +163,11 @@ def require_simulation_conflict_matches( existing: ComparisonSimulationRecord, candidate: ComparisonSimulationRecord, ) -> None: + """Raise ComparisonRunIdentityError if a duplicate insert names a different worker. + + The stored simulation and the retried record must agree on their parent, + role, input digest and worker release. + """ _require_equal_fields(existing, candidate, SIMULATION_CONFLICT_FIELDS) @@ -202,6 +214,10 @@ def require_comparison_update_only( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ComparisonRunIdentityError if anything but the comparison fields changes. + + updated_at may also differ; every other report field must match the stored row. + """ fields = tuple( field_name for field_name in ComparisonReportRecord.model_fields @@ -215,6 +231,7 @@ def require_result_comparison_unchanged( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ComparisonRunIdentityError if the result-comparison fields differ.""" _require_equal_fields(existing, candidate, REPORT_RESULT_COMPARISON_FIELDS) @@ -316,6 +333,11 @@ def require_lifecycle_transition( current: ComparisonRunLifecycleStatus, candidate: ComparisonRunLifecycleStatus, ) -> None: + """Raise ComparisonRunStateTransitionError if lifecycle cannot move to candidate. + + Every status may repeat; succeeded and skipped are final, and failed or + incomplete runs may return to running. + """ if candidate not in ALLOWED_TRANSITIONS[current]: raise ComparisonRunStateTransitionError( f"comparison-run lifecycle cannot move from {current.value} to {candidate.value}" @@ -326,6 +348,10 @@ def require_aggregation_transition( current: ComparisonRunAggregationStatus, candidate: ComparisonRunAggregationStatus, ) -> None: + """Raise ComparisonRunStateTransitionError if aggregation cannot move to candidate. + + Every status may repeat; succeeded is final and failed may return to running. + """ if candidate not in ALLOWED_AGGREGATION_TRANSITIONS[current]: raise ComparisonRunStateTransitionError( "comparison-run aggregation cannot move " @@ -337,6 +363,10 @@ def require_result_comparison_transition( current: ResultComparisonStatus, candidate: ResultComparisonStatus, ) -> None: + """Raise ComparisonRunStateTransitionError if comparison cannot move to candidate. + + Not-requested, matched, different and failed are all final. + """ if candidate not in ALLOWED_RESULT_COMPARISON_TRANSITIONS[current]: raise ComparisonRunStateTransitionError( f"result comparison cannot move from {current.value} to {candidate.value}" diff --git a/policyengine_api/services/v2/households/database_connectors/creates.py b/policyengine_api/services/v2/households/database_connectors/creates.py index 72624119c..721a0096e 100644 --- a/policyengine_api/services/v2/households/database_connectors/creates.py +++ b/policyengine_api/services/v2/households/database_connectors/creates.py @@ -21,6 +21,11 @@ def create_household( canonicalization_version: int, content_hash: str, ) -> UUID | None: + """Insert a household unless one with the same canonical version and hash exists. + + Return the new id, or None when the unique content-hash constraint made the + insert a no-op; the caller then reads the existing row. + """ household_id = uuid4() statement = ( insert(Household) @@ -50,6 +55,11 @@ def create_legacy_household_mapping( fingerprint_version: int, fingerprint_sha256: str, ) -> UUID | None: + """Insert the mapping from a v1 household id to its v2 household. + + Return the new mapping id, or None when a mapping for this country and v1 id + already exists (for example, a concurrent copy won the race). + """ mapping_id = uuid4() statement = ( insert(LegacyHouseholdMapping) diff --git a/policyengine_api/services/v2/households/database_connectors/reads.py b/policyengine_api/services/v2/households/database_connectors/reads.py index 107a72ca8..8a1e6ad8c 100644 --- a/policyengine_api/services/v2/households/database_connectors/reads.py +++ b/policyengine_api/services/v2/households/database_connectors/reads.py @@ -15,6 +15,7 @@ def read_household_by_content_identity( session: Session, *, canonicalization_version: int, content_hash: str ) -> Household | None: + """Return the household with this canonical version and content hash, or None.""" return session.exec( select(Household).where( Household.canonicalization_version == canonicalization_version, @@ -26,6 +27,7 @@ def read_household_by_content_identity( def read_household_row( session: Session, *, country_id: str, household_id: UUID ) -> Household | None: + """Return the household with this id in this country, or None.""" return session.exec( select(Household).where( Household.id == household_id, @@ -42,6 +44,10 @@ def read_household_rows( offset: int, limit: int, ) -> list[Household]: + """Return up to ``limit + 1`` households in creation order, starting at ``offset``. + + The extra row lets the caller report whether another page exists. + """ statement = select(Household).where(Household.country_id == country_id) if default_year is not None: statement = statement.where(Household.default_year == default_year) @@ -61,6 +67,11 @@ def read_legacy_household_mapping( legacy_household_id: int, lock: bool, ) -> LegacyHouseholdMapping | None: + """Return the v2 mapping for a v1 household id, or None. + + With ``lock``, the row is selected FOR UPDATE so concurrent copies of the same + v1 household serialize on it. + """ statement = select(LegacyHouseholdMapping).where( LegacyHouseholdMapping.country_id == country_id, LegacyHouseholdMapping.legacy_household_id == legacy_household_id, diff --git a/policyengine_api/services/v2/households/database_session.py b/policyengine_api/services/v2/households/database_session.py index c0d580c90..6470744dd 100644 --- a/policyengine_api/services/v2/households/database_session.py +++ b/policyengine_api/services/v2/households/database_session.py @@ -13,15 +13,20 @@ class HouseholdDatabaseSession: + """Session lifetimes for household services: read-only sessions and transactions.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/households/services.py b/policyengine_api/services/v2/households/services.py index 0bfe957fa..c698ea951 100644 --- a/policyengine_api/services/v2/households/services.py +++ b/policyengine_api/services/v2/households/services.py @@ -58,6 +58,7 @@ def normalize_creation_input( household_input: HouseholdCreationInput, ) -> HouseholdCreationInput: + """Return the input with its household document validated and normalized.""" document = household_input.household_data assert isinstance(document, dict) return HouseholdCreationInput( @@ -78,6 +79,11 @@ def create_normalized_household( [HouseholdCreationInput], CanonicalHouseholdContent ] = canonicalize_household, ) -> HouseholdCreationResult: + """Store a normalized household, or return the existing one with identical content. + + Raise HouseholdCreationIntegrityError if the hash conflict finds no row, and + HouseholdContentHashCollisionError if the stored row's canonical bytes differ. + """ content = canonicalizer(household_input) created_id = create_household( session, @@ -108,6 +114,7 @@ def create_normalized_household( def read_complete_household( session: Session, *, country_id: str, household_id: UUID ) -> HouseholdRead: + """Return one household of this country, or raise HouseholdNotFoundError.""" household = read_household_row( session, country_id=country_id, @@ -126,6 +133,7 @@ def read_household_page( offset: int, limit: int, ) -> HouseholdPage: + """Return a page of households, optionally filtered by default year.""" return household_page( read_household_rows( session, @@ -143,6 +151,12 @@ def mirror_legacy_household_in_session( session: Session, snapshot: LegacyHouseholdSnapshot, ) -> LegacyHouseholdPersistenceResult: + """Copy one committed v1 household into v2 and map its v1 id to the result. + + Replaying the same v1 household is idempotent. Raise + LegacyHouseholdMappingIntegrityError if the v1 id is already mapped with a + different source fingerprint or to a different v2 household. + """ fingerprint = legacy_household_fingerprint(snapshot) existing = read_legacy_household_mapping( session, @@ -210,12 +224,19 @@ def mirror_legacy_household_in_session( class V2HouseholdService: + """Service for creating, reading and copying immutable v2 households.""" + def __init__(self, database_session: HouseholdDatabaseSession) -> None: + """Open every session through ``database_session``.""" self._database_session = database_session def create_household( self, household_input: NativeHouseholdCreationInput ) -> NativeHouseholdCreation: + """Store native household content, deduplicated by content hash. + + Return the stored household and whether this call created it. + """ normalized = normalize_creation_input( HouseholdCreationInput.model_validate(household_input.model_dump()) ) @@ -229,6 +250,7 @@ def create_household( return NativeHouseholdCreation(item=item, created=result.created) def get_household(self, *, country_id: str, household_id: UUID) -> HouseholdRead: + """Return one household of this country, or raise HouseholdNotFoundError.""" with self._database_session.read() as session: return read_complete_household( session, @@ -244,6 +266,7 @@ def list_households( offset: int, limit: int, ) -> HouseholdPage: + """Return a page of households, optionally filtered by default year.""" with self._database_session.read() as session: return read_household_page( session, @@ -256,5 +279,6 @@ def list_households( def mirror_legacy_household( self, snapshot: LegacyHouseholdSnapshot ) -> LegacyHouseholdPersistenceResult: + """Copy one committed v1 household into v2 in its own transaction.""" with self._database_session.transaction() as session: return mirror_legacy_household_in_session(session, snapshot) diff --git a/policyengine_api/services/v2/households/transformations.py b/policyengine_api/services/v2/households/transformations.py index def489d9d..a8cf8127f 100644 --- a/policyengine_api/services/v2/households/transformations.py +++ b/policyengine_api/services/v2/households/transformations.py @@ -72,6 +72,11 @@ def _canonical_json( def canonical_household_document(household_input: HouseholdCreationInput) -> bytes: + """Serialize a household to the canonical bytes that its content hash covers. + + People and entity records are sorted by id and keys are sorted, so the same + household always yields the same bytes regardless of input order. + """ content = dict(cast("dict[str, Any]", household_input.household_data)) for collection, records in content.items(): if collection in { @@ -96,6 +101,7 @@ def canonical_household_document(household_input: HouseholdCreationInput) -> byt def canonicalize_household( household_input: HouseholdCreationInput, ) -> CanonicalHouseholdContent: + """Return the canonical bytes, SHA-256 hash and canonicalization version.""" document = canonical_household_document(household_input) return CanonicalHouseholdContent( version=HOUSEHOLD_CANONICALIZATION_VERSION, @@ -105,6 +111,7 @@ def canonicalize_household( def stored_household_creation_input(household: Household) -> HouseholdCreationInput: + """Rebuild a stored row's creation input so its canonical bytes can be compared.""" return HouseholdCreationInput.model_validate( { "country_id": household.country_id, @@ -115,6 +122,7 @@ def stored_household_creation_input(household: Household) -> HouseholdCreationIn def household_read(household: Household) -> HouseholdRead: + """Convert a stored household row to the service's read model.""" return HouseholdRead( id=household.id, country_id=household.country_id, @@ -126,6 +134,10 @@ def household_read(household: Household) -> HouseholdRead: def household_page(rows: list[Household], *, offset: int, limit: int) -> HouseholdPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ displayed = rows[:limit] return HouseholdPage( items=tuple(household_read(row) for row in displayed), @@ -176,6 +188,13 @@ def _derive_default_year(values: list[object]) -> int | None: def translate_legacy_household( snapshot: LegacyHouseholdSnapshot, ) -> HouseholdCreationInput: + """Translate a v1 household document into normalized v2 household content. + + Named v1 people and groups get stable ids in name order, and group members + become person memberships. The default year is derived from the value periods. + Raise LegacyHouseholdTranslationError when the v1 document cannot be + translated exactly. + """ source = snapshot.household_json assert isinstance(source, dict) collections = SUPPORTED_ENTITY_COLLECTIONS[snapshot.country_id] @@ -277,6 +296,7 @@ def translate_legacy_household( def legacy_household_fingerprint(snapshot: LegacyHouseholdSnapshot) -> str: + """Return a v1 snapshot's SHA-256 fingerprint, which detects changed v1 content.""" document = { "fingerprint_version": LEGACY_HOUSEHOLD_FINGERPRINT_VERSION, **snapshot.model_dump(mode="json"), diff --git a/policyengine_api/services/v2/households/types.py b/policyengine_api/services/v2/households/types.py index dab71705a..23eefa1cb 100644 --- a/policyengine_api/services/v2/households/types.py +++ b/policyengine_api/services/v2/households/types.py @@ -40,6 +40,7 @@ class HouseholdCreationInput(StrictHouseholdInput): @field_validator("household_data") @classmethod def require_document(cls, value: object) -> object: + """Reject household_data that is not a JSON object.""" if type(value) is not dict: raise ValueError("household_data must be an object") return value @@ -64,6 +65,7 @@ class LegacyHouseholdSnapshot(StrictHouseholdInput): @field_validator("household_json") @classmethod def require_legacy_document(cls, value: object) -> object: + """Reject a v1 household_json that is not a JSON object.""" if type(value) is not dict: raise ValueError("legacy household_json must be an object") return value @@ -71,6 +73,12 @@ def require_legacy_document(cls, value: object) -> object: @dataclass(frozen=True) class CanonicalHouseholdContent: + """Canonical bytes of one household and the SHA-256 digest that deduplicates it. + + ``version`` is the canonicalization version; the same content hashed under + a different version is a different stored household. + """ + version: int document: bytes content_hash: str @@ -78,12 +86,20 @@ class CanonicalHouseholdContent: @dataclass(frozen=True) class HouseholdCreationResult: + """The stored household id, and whether this call inserted it or found it.""" + household_id: UUID created: bool @dataclass(frozen=True) class LegacyHouseholdPersistenceResult: + """Outcome of copying one v1 household into v2. + + ``household_id`` is the v2 household the v1 id maps to; the flags say whether + this call created the household and the v1-to-v2 mapping or found them. + """ + household_id: UUID household_created: bool mapping_created: bool @@ -91,6 +107,8 @@ class LegacyHouseholdPersistenceResult: @dataclass(frozen=True) class HouseholdRead: + """A stored v2 household as the service returns it to routes.""" + id: UUID country_id: str default_year: int | None @@ -101,6 +119,8 @@ class HouseholdRead: @dataclass(frozen=True) class HouseholdPage: + """One page of households in creation order; ``has_more`` says if more follow.""" + items: tuple[HouseholdRead, ...] offset: int limit: int @@ -109,5 +129,7 @@ class HouseholdPage: @dataclass(frozen=True) class NativeHouseholdCreation: + """A created-or-found household, with ``created`` choosing 201 or 200.""" + item: HouseholdRead created: bool diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads.py b/policyengine_api/services/v2/metadata/database_connectors/reads.py index 48690c061..bd91984b2 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads.py @@ -23,6 +23,11 @@ def read_metadata_catalog( running_policyengine_version: str, policyengine_version: str | None, ) -> SelectedCatalog: + """Return the country's catalog for policyengine_version, else the running version. + + Raise the catalog_selection errors for an unsupported country, an + unpublished version or an unreachable database. + """ return select_catalog( session, country_id=country_id, @@ -35,6 +40,10 @@ def read_rows( session: Session, statement: Any, # noqa: ANN401 # Session.exec overloads omit core SQLAlchemy selects ) -> list[Any]: + """Execute statement and return every row. + + Raise MetadataCatalogUnavailableError when the database query fails. + """ try: return list(session.exec(statement).all()) except SQLAlchemyError as error: diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_datasets.py b/policyengine_api/services/v2/metadata/database_connectors/reads_datasets.py index 27a197bae..cd5dbc276 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_datasets.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_datasets.py @@ -16,6 +16,10 @@ def read_datasets( session: Session, *, model_version_id: UUID, offset: int, limit: int ) -> list[Dataset]: + """Return a page of a model version's logical input datasets in name order. + + Fetch limit + 1 rows so the caller can tell whether another page exists. + """ return read_rows( session, select(Dataset) @@ -33,6 +37,7 @@ def read_datasets( def read_dataset( session: Session, *, model_version_id: UUID, dataset_id: UUID ) -> Dataset | None: + """Return the model version's logical input dataset with dataset_id, or None.""" rows = read_rows( session, select(Dataset).where( diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py b/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py index 230793152..d29390ef5 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py @@ -150,6 +150,13 @@ def read_parameter_children( offset: int, limit: int, ) -> list[Any]: + """Return a page of the direct children of parent_path in a model version's tree. + + Each row has the child's path, label, type (node or parameter), child count + and, for parameters, the parameter's summary columns. dialect selects SQLite + or PostgreSQL string functions; limit + 1 rows are fetched so the caller can + tell whether another page exists. + """ return read_rows( session, _parameter_children_statement( diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_parameters.py b/policyengine_api/services/v2/metadata/database_connectors/reads_parameters.py index 848804c66..85fa4dc0d 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_parameters.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_parameters.py @@ -27,6 +27,12 @@ def read_parameters( limit: int, search: str | None, ) -> list[Parameter]: + """Return a page of a model version's parameters in name order. + + search matches name, label or description as a literal, case-insensitive + substring. Fetch limit + 1 rows so the caller can tell whether another page + exists. + """ statement = select(Parameter).where( col(Parameter.tax_benefit_model_version_id) == model_version_id ) @@ -48,6 +54,7 @@ def read_parameters( def read_parameter( session: Session, *, model_version_id: UUID, parameter_id: UUID ) -> Parameter | None: + """Return the model version's parameter with parameter_id, or None.""" rows = read_rows( session, select(Parameter).where( @@ -67,6 +74,13 @@ def read_parameter_values( offset: int, limit: int, ) -> list[ParameterValue]: + """Return a page of a model version's current-law parameter values. + + Values tied to a policy or dynamic are excluded. parameter_id keeps one + parameter's values and selected_day keeps those in effect on that day. Rows + are ordered by parameter name, then newest first, and limit + 1 are fetched + so the caller can tell whether another page exists. + """ statement = ( select(ParameterValue) .join(Parameter, col(Parameter.id) == col(ParameterValue.parameter_id)) @@ -101,6 +115,7 @@ def read_parameter_values( def read_parameter_value( session: Session, *, model_version_id: UUID, value_id: UUID ) -> ParameterValue | None: + """Return the model version's current-law value with value_id, or None.""" rows = read_rows( session, select(ParameterValue) diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_regions.py b/policyengine_api/services/v2/metadata/database_connectors/reads_regions.py index 719f77a20..e7dacf1a9 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_regions.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_regions.py @@ -21,6 +21,11 @@ def read_regions( offset: int | None = None, limit: int | None = None, ) -> list[Region]: + """Return a model version's regions in code order, optionally of one type. + + Without a limit every region is returned; with one, limit + 1 rows are + fetched so the caller can tell whether another page exists. + """ statement = select(Region).where( col(Region.tax_benefit_model_version_id) == model_version_id ) @@ -41,6 +46,7 @@ def read_region( region_id: UUID | None = None, region_code: str | None = None, ) -> Region | None: + """Return the model version's region matching the given id and code, or None.""" statement = select(Region).where( col(Region.tax_benefit_model_version_id) == model_version_id ) @@ -55,6 +61,10 @@ def read_region( def read_input_dataset( session: Session, *, model_version_id: UUID, dataset_id: UUID ) -> Dataset | None: + """Return the model version's logical input dataset with dataset_id, or None. + + Used to load a region's default dataset for the economy options. + """ rows = read_rows( session, select(Dataset).where( diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_variables.py b/policyengine_api/services/v2/metadata/database_connectors/reads_variables.py index f833c2257..4d84be396 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_variables.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_variables.py @@ -26,6 +26,12 @@ def read_variables( limit: int, search: str | None, ) -> list[Variable]: + """Return a page of a model version's variables in name order. + + search matches name, label or description as a literal, case-insensitive + substring. Fetch limit + 1 rows so the caller can tell whether another page + exists. + """ statement = select(Variable).where( col(Variable.tax_benefit_model_version_id) == model_version_id ) @@ -47,6 +53,7 @@ def read_variables( def read_variable( session: Session, *, model_version_id: UUID, variable_id: UUID ) -> Variable | None: + """Return the model version's variable with variable_id, or None.""" rows = read_rows( session, select(Variable).where( diff --git a/policyengine_api/services/v2/metadata/database_session.py b/policyengine_api/services/v2/metadata/database_session.py index e0e1fb278..207b7d07e 100644 --- a/policyengine_api/services/v2/metadata/database_session.py +++ b/policyengine_api/services/v2/metadata/database_session.py @@ -12,15 +12,19 @@ class MetadataDatabaseSession: """Own the request-scoped read session used by metadata services.""" def __init__(self, session: Session) -> None: + """Wrap one request's session; the service reads through it until close.""" self._session = session @property def session(self) -> Session: + """The open request session that every metadata query runs on.""" return self._session @property def dialect_name(self) -> str: + """The bound dialect's name; it picks SQLite or PostgreSQL string functions.""" return self._session.get_bind().dialect.name def close(self) -> None: + """Close the request's session.""" self._session.close() diff --git a/policyengine_api/services/v2/metadata/services.py b/policyengine_api/services/v2/metadata/services.py index 948743543..8082c6fad 100644 --- a/policyengine_api/services/v2/metadata/services.py +++ b/policyengine_api/services/v2/metadata/services.py @@ -77,7 +77,13 @@ class V2MetadataService: - """Sequence metadata reads through a request-scoped database session.""" + """Sequence metadata reads through a request-scoped database session. + + Each read selects the country's catalog for the requested PolicyEngine.py + version, or for the running version when none is given. List methods page + with offset and a limit of 1-500 and report has_more; get methods raise + MetadataResourceNotFoundError when the selected catalog lacks the item. + """ def __init__( self, @@ -85,12 +91,14 @@ def __init__( *, running_policyengine_version: str, ) -> None: + """Bind to a request's session and the running PolicyEngine.py version.""" self._database_session = database_session self._running_policyengine_version = validate_policyengine_version( running_policyengine_version ) def close(self) -> None: + """Close the request's database session.""" self._database_session.close() def _select_catalog( @@ -122,6 +130,7 @@ def list_models( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataModel]: + """Return the catalog's only model on the first page; later pages are empty.""" selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -134,6 +143,7 @@ def get_model( model_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataModel]: + """Return the selected catalog's model if its id is model_id.""" selected = self._select_catalog(country_id, policyengine_version) item = require_metadata_resource( metadata_model(selected) if selected.model.id == model_id else None, @@ -146,6 +156,7 @@ def get_model( def get_model_by_country( self, country_id: str, policyengine_version: str | None = None ) -> MetadataModelSelectionResult: + """Return the country's model together with the selected model version.""" return metadata_model_selection( self._select_catalog(country_id, policyengine_version) ) @@ -158,6 +169,7 @@ def list_model_versions( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataModelVersionDetail]: + """Return the selected version on the first page; later pages are empty.""" selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -170,6 +182,7 @@ def get_model_version( version_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataModelVersionDetail]: + """Return the selected model version if its id is version_id.""" selected = self._select_catalog(country_id, policyengine_version) item = require_metadata_resource( ( @@ -192,6 +205,11 @@ def list_variables( limit: int = 100, search: str | None = None, ) -> MetadataPageResult[MetadataVariable]: + """Return a page of the catalog's variables in name order. + + search keeps variables whose name, label or description contains it, + ignoring case. + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -215,6 +233,7 @@ def get_variable( variable_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataVariable]: + """Return the variable with variable_id if it is in the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_variable( @@ -238,6 +257,11 @@ def list_parameters( limit: int = 100, search: str | None = None, ) -> MetadataPageResult[MetadataParameterSummary]: + """Return a page of the catalog's parameters in name order. + + search keeps parameters whose name, label or description contains it, + ignoring case. + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -261,6 +285,7 @@ def get_parameter( parameter_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataParameterSummary]: + """Return the parameter with parameter_id if it is in the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_parameter( @@ -284,6 +309,11 @@ def list_parameter_children( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataParameterChild]: + """Return a page of the direct children of parent_path in the parameter tree. + + An empty parent_path lists the top-level children. A child is a parameter + when it is a leaf with a parameter row, and a node otherwise. + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -313,6 +343,12 @@ def list_parameter_values( limit: int = 100, now: datetime | None = None, ) -> MetadataPageResult[MetadataCanonicalParameterValue]: + """Return a page of current-law parameter values, newest first per parameter. + + Values belonging to a reform or dynamic are never listed. parameter_id keeps + one parameter's values, and current=True keeps those in effect on the UTC + day of now (today by default). + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -338,6 +374,10 @@ def get_parameter_value( value_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataCanonicalParameterValue]: + """Return the current-law value with value_id if it is in the selected catalog. + + Values belonging to a reform or dynamic are never returned. + """ selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_parameter_value( @@ -360,6 +400,10 @@ def list_datasets( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataDataset]: + """Return a page of the catalog's logical input datasets in name order. + + Output datasets and datasets with a storage path are never listed. + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -382,6 +426,7 @@ def get_dataset( dataset_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataDataset]: + """Return the logical input dataset with dataset_id in the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_dataset( @@ -405,6 +450,7 @@ def list_regions( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataRegion]: + """Return a page of the catalog's regions by code, optionally of one type.""" selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -428,6 +474,7 @@ def get_region( region_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataRegion]: + """Return the region with region_id from the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_region( @@ -448,6 +495,7 @@ def get_region_by_code( region_code: str, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataRegion]: + """Return the region with region_code from the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_region( @@ -465,6 +513,11 @@ def get_region_by_code( def get_economy_options( self, country_id: str, policyengine_version: str | None = None ) -> MetadataEconomyOptionsResult: + """Return the region, year and dataset choices for an economy calculation. + + Raise MetadataCatalogUnavailableError when the catalog lacks the national + region, its default input dataset, the current-law id or the years. + """ selected = self._select_catalog(country_id, policyengine_version) regions = read_regions( self._database_session.session, diff --git a/policyengine_api/services/v2/metadata/transformations.py b/policyengine_api/services/v2/metadata/transformations.py index e17a9fd57..588869af4 100644 --- a/policyengine_api/services/v2/metadata/transformations.py +++ b/policyengine_api/services/v2/metadata/transformations.py @@ -44,6 +44,11 @@ def page_result( offset: int, limit: int, ) -> MetadataPageResult[ResourceT]: + """Wrap fetched rows as one page of the selected catalog. + + rows holds up to limit + 1 items; an extra item only sets has_more and is + dropped from the page. + """ return MetadataPageResult( policyengine_version=selected.policyengine_version, items=rows[:limit], @@ -54,6 +59,7 @@ def page_result( def metadata_model(selected: SelectedCatalog) -> MetadataModel: + """Describe the catalog's model, using the selected version's description.""" return MetadataModel( id=selected.model.id, name=selected.model.name, @@ -62,6 +68,7 @@ def metadata_model(selected: SelectedCatalog) -> MetadataModel: def metadata_model_version(selected: SelectedCatalog) -> MetadataModelVersionDetail: + """Describe the selected model version, with its current-law id and years.""" return MetadataModelVersionDetail( id=selected.model_version.id, model_id=selected.model.id, @@ -73,6 +80,7 @@ def metadata_model_version(selected: SelectedCatalog) -> MetadataModelVersionDet def metadata_model_selection(selected: SelectedCatalog) -> MetadataModelSelectionResult: + """Pair the catalog's model with its selected version.""" return MetadataModelSelectionResult( policyengine_version=selected.policyengine_version, model=metadata_model(selected), @@ -81,6 +89,7 @@ def metadata_model_selection(selected: SelectedCatalog) -> MetadataModelSelectio def metadata_variable(variable: Variable) -> MetadataVariable: + """Return the API representation of a stored variable.""" return MetadataVariable( id=variable.id, name=variable.name, @@ -96,6 +105,7 @@ def metadata_variable(variable: Variable) -> MetadataVariable: def metadata_parameter(parameter: Parameter) -> MetadataParameterSummary: + """Return the API representation of a stored parameter.""" return MetadataParameterSummary( id=parameter.id, name=parameter.name, @@ -107,6 +117,7 @@ def metadata_parameter(parameter: Parameter) -> MetadataParameterSummary: def metadata_parameter_value(value: ParameterValue) -> MetadataCanonicalParameterValue: + """Return the API representation of a stored value, exposing value_json as value.""" return MetadataCanonicalParameterValue( id=value.id, parameter_id=value.parameter_id, @@ -117,6 +128,7 @@ def metadata_parameter_value(value: ParameterValue) -> MetadataCanonicalParamete def utc_day_start(selected_time: datetime) -> datetime: + """Return midnight UTC on the day of selected_time, treating a naive time as UTC.""" if selected_time.tzinfo is None: selected_time = selected_time.replace(tzinfo=UTC) return selected_time.astimezone(UTC).replace( @@ -125,6 +137,7 @@ def utc_day_start(selected_time: datetime) -> datetime: def metadata_dataset(dataset: Dataset) -> MetadataDataset: + """Return the API representation of a stored logical input dataset.""" return MetadataDataset( id=dataset.id, name=dataset.name, @@ -134,6 +147,7 @@ def metadata_dataset(dataset: Dataset) -> MetadataDataset: def metadata_region(region: Region) -> MetadataRegion: + """Return the API representation of a stored region.""" return MetadataRegion( id=region.id, code=region.code, @@ -151,6 +165,11 @@ def metadata_region(region: Region) -> MetadataRegion: def metadata_parameter_children(rows: list[Any]) -> list[MetadataParameterChild]: + """Build parameter-tree children from read_parameter_children rows. + + Parameter leaves carry the parameter's summary, and a child with no label is + labelled with the last segment of its path. + """ items = [] for row in rows: parameter = None @@ -178,6 +197,10 @@ def metadata_parameter_children(rows: list[Any]) -> list[MetadataParameterChild] def metadata_economy_options( selected: SelectedCatalog, *, regions: list[Region], national_dataset: Dataset ) -> MetadataEconomyOptionsResult: + """Build economy options from every region, each year and the national dataset. + + The national dataset is the only dataset choice, and it is the default. + """ return MetadataEconomyOptionsResult( policyengine_version=selected.policyengine_version, current_law_id=selected.model_version.current_law_id, diff --git a/policyengine_api/services/v2/metadata/types.py b/policyengine_api/services/v2/metadata/types.py index fef5b7894..b74333291 100644 --- a/policyengine_api/services/v2/metadata/types.py +++ b/policyengine_api/services/v2/metadata/types.py @@ -11,10 +11,14 @@ class StrictResponseModel(BaseModel): + """Base for metadata results that rejects unknown fields and non-finite numbers.""" + model_config = ConfigDict(extra="forbid", allow_inf_nan=False) class MetadataRegionType(StrEnum): + """The geographic level of a region.""" + NATIONAL = "national" COUNTRY = "country" STATE = "state" @@ -26,12 +30,16 @@ class MetadataRegionType(StrEnum): class MetadataModel(StrictResponseModel): + """A country's tax-benefit model in the selected PolicyEngine.py catalog.""" + id: UUID name: str description: str | None class MetadataModelVersion(StrictResponseModel): + """One published version of a country's tax-benefit model.""" + id: UUID model_id: UUID version: str @@ -39,6 +47,12 @@ class MetadataModelVersion(StrictResponseModel): class MetadataVariable(StrictResponseModel): + """A variable the tax-benefit model computes or accepts as input. + + adds and subtracts, when present, list the variables summed into or + subtracted from this one. + """ + id: UUID name: str label: str | None @@ -52,6 +66,11 @@ class MetadataVariable(StrictResponseModel): class MetadataParameterSummary(StrictResponseModel): + """A policy parameter, described without its values. + + Its current-law values are listed by GET /v2/parameter-values. + """ + id: UUID name: str label: str | None @@ -61,6 +80,12 @@ class MetadataParameterSummary(StrictResponseModel): class MetadataCanonicalParameterValue(StrictResponseModel): + """One current-law value of a parameter. + + The value applies from start_date until end_date, or indefinitely when + end_date is null. + """ + id: UUID parameter_id: UUID value: JsonValue @@ -69,6 +94,12 @@ class MetadataCanonicalParameterValue(StrictResponseModel): class MetadataDataset(StrictResponseModel): + """A named input dataset available to the selected model version. + + These datasets are logical inputs, so storage_path is always null and + is_output_dataset is always false. + """ + id: UUID name: str description: str | None @@ -78,6 +109,12 @@ class MetadataDataset(StrictResponseModel): class MetadataRegion(StrictResponseModel): + """A region the model can simulate, such as a country, state or constituency. + + When requires_filter is true, the region is simulated by filtering its + default dataset on filter_field and filter_value. + """ + id: UUID code: str label: str @@ -93,28 +130,41 @@ class MetadataRegion(StrictResponseModel): class MetadataRegionOption(StrictResponseModel): + """A region that can be chosen for an economy calculation; name is its code.""" + name: str label: str type: MetadataRegionType class MetadataTimePeriodOption(StrictResponseModel): + """A year that can be chosen for an economy calculation.""" + name: int label: str class MetadataDatasetOption(StrictResponseModel): + """The national dataset, offered as the default for an economy calculation.""" + name: str label: str default: Literal[True] = True class MetadataModelVersionDetail(MetadataModelVersion): + """A model version, with its current-law policy id and the years offered.""" + current_law_id: int metadata_time_periods: list[int] class MetadataParameterChild(StrictResponseModel): + """A direct child in the parameter tree. + + A node has a child_count; a parameter leaf carries the parameter's summary. + """ + path: str label: str type: Literal["node", "parameter"] @@ -126,6 +176,8 @@ class MetadataParameterChild(StrictResponseModel): class MetadataPageResult(StrictResponseModel, Generic[ResourceT]): + """One page of a metadata collection; has_more says whether a later page exists.""" + policyengine_version: str items: list[ResourceT] offset: int @@ -134,17 +186,27 @@ class MetadataPageResult(StrictResponseModel, Generic[ResourceT]): class MetadataDetailResult(StrictResponseModel, Generic[ResourceT]): + """One metadata item, with the PolicyEngine.py version it was read from.""" + policyengine_version: str item: ResourceT class MetadataModelSelectionResult(StrictResponseModel): + """A country's model together with the model version selected for the request.""" + policyengine_version: str model: MetadataModel model_version: MetadataModelVersionDetail class MetadataEconomyOptionsResult(StrictResponseModel): + """The choices for configuring an economy calculation. + + It lists the selectable regions, years and datasets, and the id of the + current-law policy to use as a baseline. + """ + policyengine_version: str current_law_id: int region: list[MetadataRegionOption] diff --git a/policyengine_api/services/v2/metadata/validators.py b/policyengine_api/services/v2/metadata/validators.py index 45bb3c88f..036722af4 100644 --- a/policyengine_api/services/v2/metadata/validators.py +++ b/policyengine_api/services/v2/metadata/validators.py @@ -22,6 +22,10 @@ class InvalidMetadataPageError(ValueError): def validate_metadata_page(offset: int, limit: int) -> tuple[int, int]: + """Return offset and limit unchanged, or raise InvalidMetadataPageError. + + offset must be at least 0 and limit between 1 and 500. + """ if offset < 0: raise InvalidMetadataPageError("offset must be at least 0") if not 1 <= limit <= 500: @@ -35,6 +39,7 @@ def validate_metadata_page(offset: int, limit: int) -> tuple[int, int]: def require_metadata_resource( resource: ResourceT | None, *, description: str ) -> ResourceT: + """Return resource, or raise MetadataResourceNotFoundError naming description.""" if resource is None: raise MetadataResourceNotFoundError(f"{description} was not found") return resource diff --git a/policyengine_api/services/v2/policies/database_connectors/creates.py b/policyengine_api/services/v2/policies/database_connectors/creates.py index ccebd0bb3..18e81b853 100644 --- a/policyengine_api/services/v2/policies/database_connectors/creates.py +++ b/policyengine_api/services/v2/policies/database_connectors/creates.py @@ -21,6 +21,12 @@ def create_policy( canonicalization_version: int, content_hash: str, ) -> UUID | None: + """Insert a policy unless one with the same canonical version and hash exists. + + Return the new id, or None when the unique content-hash constraint made the + insert a no-op; the caller then reads the existing row. Parameter values are + inserted separately by create_parameter_values. + """ policy_id = uuid4() statement = ( insert(Policy) @@ -44,6 +50,7 @@ def create_parameter_values( policy_id: UUID, policy_input: ResolvedPolicyCreationInput, ) -> None: + """Insert a new policy's parameter values and flush them so later reads see them.""" session.add_all( [ ParameterValue( @@ -68,6 +75,11 @@ def create_legacy_policy_mapping( source_policy_hash: str, policy_id: UUID, ) -> UUID | None: + """Insert the mapping from a v1 policy id to its v2 policy. + + Return the new mapping id, or None when a mapping for this country and v1 id + already exists (for example, a concurrent copy won the race). + """ return session.execute( insert(LegacyPolicyMapping) .values( diff --git a/policyengine_api/services/v2/policies/database_connectors/reads.py b/policyengine_api/services/v2/policies/database_connectors/reads.py index 244bb9eeb..6c649b384 100644 --- a/policyengine_api/services/v2/policies/database_connectors/reads.py +++ b/policyengine_api/services/v2/policies/database_connectors/reads.py @@ -30,6 +30,10 @@ def read_policy_catalog( policyengine_version: str | None = None, running_policyengine_version: str = POLICYENGINE_VERSION, ) -> SelectedCatalog: + """Select the country's metadata catalog for the requested PolicyEngine version. + + Without ``policyengine_version``, select the catalog of the running version. + """ return select_catalog( session, country_id=country_id, @@ -41,6 +45,7 @@ def read_policy_catalog( def read_version_parameter_ids( session: Session, *, model_version_id: UUID, requested_ids: set[UUID] ) -> set[UUID]: + """Return the requested parameter ids that exist in this model version.""" if not requested_ids: return set() return set( @@ -56,6 +61,7 @@ def read_version_parameter_ids( def read_parameters_by_name( session: Session, *, model_version_id: UUID, names: set[str] ) -> dict[str, Parameter]: + """Return this model version's parameters with the given names, keyed by name.""" if not names: return {} parameters = session.exec( @@ -70,6 +76,7 @@ def read_parameters_by_name( def read_policy_by_content_identity( session: Session, *, canonicalization_version: int, content_hash: str ) -> Policy | None: + """Return the policy with this canonical version and content hash, or None.""" return session.exec( select(Policy).where( Policy.canonicalization_version == canonicalization_version, @@ -81,10 +88,12 @@ def read_policy_by_content_identity( def read_model_version( session: Session, model_version_id: UUID ) -> TaxBenefitModelVersion | None: + """Return the model version with this id, or None.""" return session.get(TaxBenefitModelVersion, model_version_id) def read_parameter_values(session: Session, policy_id: UUID) -> list[ParameterValue]: + """Return every parameter value stored for one policy, in no particular order.""" return list( session.exec( select(ParameterValue).where(ParameterValue.policy_id == policy_id) @@ -95,6 +104,11 @@ def read_parameter_values(session: Session, policy_id: UUID) -> list[ParameterVa def read_legacy_policy_mapping( session: Session, *, country_id: str, legacy_policy_id: int, lock: bool ) -> LegacyPolicyMapping | None: + """Return the v2 mapping for a v1 policy id, or None. + + With ``lock``, the row is selected FOR UPDATE so concurrent copies of the same + v1 policy serialize on it. + """ statement = select(LegacyPolicyMapping).where( LegacyPolicyMapping.country_id == country_id, LegacyPolicyMapping.legacy_policy_id == legacy_policy_id, @@ -107,6 +121,7 @@ def read_legacy_policy_mapping( def read_policy_row( session: Session, *, country_id: str, policy_id: UUID ) -> Policy | None: + """Return the policy with this id in this country, or None.""" return session.exec( select(Policy).where(Policy.id == policy_id, Policy.country_id == country_id) ).one_or_none() @@ -120,6 +135,10 @@ def read_policy_rows( offset: int, limit: int, ) -> list[Policy]: + """Return up to ``limit + 1`` policies in creation order, starting at ``offset``. + + The extra row lets the caller report whether another page exists. + """ statement = select(Policy).where(Policy.country_id == country_id) if tax_benefit_model_id is not None: statement = statement.where(Policy.tax_benefit_model_id == tax_benefit_model_id) @@ -135,6 +154,7 @@ def read_policy_rows( def read_parameter_values_with_names( session: Session, policy_ids: list[UUID] ) -> list[tuple[ParameterValue, str]]: + """Return the policies' values with parameter names, by name then start date.""" if not policy_ids: return [] return list( diff --git a/policyengine_api/services/v2/policies/database_session.py b/policyengine_api/services/v2/policies/database_session.py index af57e150e..dbc02bdab 100644 --- a/policyengine_api/services/v2/policies/database_session.py +++ b/policyengine_api/services/v2/policies/database_session.py @@ -16,14 +16,17 @@ class PolicyDatabaseSession: """Open read sessions and transaction-scoped sessions for policy services.""" def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/policies/services.py b/policyengine_api/services/v2/policies/services.py index 0fcb68c23..9a8e0d3b2 100644 --- a/policyengine_api/services/v2/policies/services.py +++ b/policyengine_api/services/v2/policies/services.py @@ -70,6 +70,11 @@ def resolve_policy_creation_input( policyengine_version: str | None = None, running_policyengine_version: str = POLICYENGINE_VERSION, ) -> ResolvedPolicyCreationInput: + """Bind policy content to the selected catalog's model version. + + Raise PolicyCatalogValidationError if the model or any parameter id is not + part of the selected catalog. + """ selected = read_policy_catalog( session, policy_input.country_id, @@ -99,6 +104,11 @@ def create_resolved_policy( [ResolvedPolicyCreationInput], CanonicalPolicyContent ] = canonicalize_policy, ) -> PolicyCreationResult: + """Store a resolved policy, or return the existing one with identical content. + + Raise PolicyCreationIntegrityError if the hash conflict finds no usable row, + and PolicyContentHashCollisionError if the stored row's canonical bytes differ. + """ content = canonicalizer(policy_input) created_policy_id = create_policy( session, @@ -143,6 +153,7 @@ def create_resolved_policy( def read_complete_policy( session: Session, *, country_id: str, policy_id: UUID ) -> PolicyRead: + """Return one policy of this country, or raise PolicyNotFoundError.""" policy = read_policy_row(session, country_id=country_id, policy_id=policy_id) if policy is None: raise PolicyNotFoundError(f"policy {policy_id} was not found") @@ -160,6 +171,7 @@ def read_policy_page( offset: int, limit: int, ) -> PolicyPage: + """Return a page of policies with their values, optionally for one model.""" rows = read_policy_rows( session, country_id=country_id, @@ -182,6 +194,12 @@ def mirror_legacy_policy_in_session( running_policyengine_version: str = POLICYENGINE_VERSION, country_package_versions: Mapping[str, str] = COUNTRY_PACKAGE_VERSIONS, ) -> LegacyPolicyPersistenceResult: + """Copy one committed v1 policy into v2 and map its v1 id to the result. + + The policy is translated against the running catalog. Replaying the same v1 + policy is idempotent. Raise LegacyPolicyMappingIntegrityError if the v1 id is + already mapped with a different source hash or to a different v2 policy. + """ existing = read_legacy_policy_mapping( session, country_id=snapshot.country_id, @@ -268,12 +286,20 @@ def __init__( *, running_policyengine_version: str = POLICYENGINE_VERSION, ) -> None: + """Open sessions through ``database_session``. + + ``running_policyengine_version`` selects the catalog when a request names none. + """ self._database_session = database_session self._running_policyengine_version = running_policyengine_version def create_policy( self, policy_input: NativePolicyCreationInput ) -> NativePolicyCreation: + """Store native policy content, deduplicated by content hash. + + Return the stored policy and whether this call created it. + """ content = PolicyCreationInput.model_validate( policy_input.model_dump(exclude={"policyengine_version"}) ) @@ -293,6 +319,7 @@ def create_policy( return NativePolicyCreation(item=item, created=persisted.created) def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: + """Return one policy of this country, or raise PolicyNotFoundError.""" with self._database_session.read() as session: return read_complete_policy( session, country_id=country_id, policy_id=policy_id @@ -306,6 +333,7 @@ def list_policies( offset: int = 0, limit: int = 100, ) -> PolicyPage: + """Return a page of policies, optionally for one tax-benefit model.""" with self._database_session.read() as session: return read_policy_page( session, @@ -318,6 +346,7 @@ def list_policies( def mirror_legacy_policy( self, snapshot: LegacyPolicySnapshot ) -> LegacyPolicyPersistenceResult: + """Copy one committed v1 policy into v2 in its own transaction.""" with self._database_session.transaction() as session: return mirror_legacy_policy_in_session( session, diff --git a/policyengine_api/services/v2/policies/transformations.py b/policyengine_api/services/v2/policies/transformations.py index f80e04390..ef124f337 100644 --- a/policyengine_api/services/v2/policies/transformations.py +++ b/policyengine_api/services/v2/policies/transformations.py @@ -75,11 +75,17 @@ def _canonical_json( def canonical_utc_datetime(value: datetime) -> str: + """Render a datetime in UTC with microseconds and a ``Z`` suffix, as hashed.""" utc_value = value.astimezone(UTC) return utc_value.isoformat(timespec="microseconds").replace("+00:00", "Z") def canonical_policy_document(policy_input: ResolvedPolicyCreationInput) -> bytes: + """Serialize a resolved policy to the canonical bytes that its content hash covers. + + Values are sorted by parameter id and dates, and keys are sorted, so the same + policy always yields the same bytes regardless of input order. + """ parameter_values = sorted( policy_input.parameter_values, key=lambda value: ( @@ -113,6 +119,7 @@ def canonical_policy_document(policy_input: ResolvedPolicyCreationInput) -> byte def canonicalize_policy( policy_input: ResolvedPolicyCreationInput, ) -> CanonicalPolicyContent: + """Return the canonical bytes, SHA-256 hash and canonicalization version.""" document = canonical_policy_document(policy_input) return CanonicalPolicyContent( version=POLICY_CANONICALIZATION_VERSION, @@ -126,6 +133,7 @@ def stored_policy_creation_input( model_version: TaxBenefitModelVersion, values: list[ParameterValue], ) -> ResolvedPolicyCreationInput: + """Rebuild a stored policy's input so its canonical bytes can be compared.""" return ResolvedPolicyCreationInput.model_validate( { "country_id": policy.country_id, @@ -149,6 +157,7 @@ def policy_parameter_values_by_policy( policy_ids: list[UUID], rows: list[tuple[ParameterValue, str]], ) -> dict[UUID, tuple[PolicyParameterValueRead, ...]]: + """Group named value rows by policy id; every requested policy gets an entry.""" grouped: dict[UUID, list[PolicyParameterValueRead]] = { policy_id: [] for policy_id in policy_ids } @@ -172,6 +181,7 @@ def policy_read( policy: Policy, values: dict[UUID, tuple[PolicyParameterValueRead, ...]], ) -> PolicyRead: + """Convert a stored policy and its grouped values to the service's read model.""" return PolicyRead( id=policy.id, country_id=policy.country_id, @@ -190,6 +200,10 @@ def policy_page( offset: int, limit: int, ) -> PolicyPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ policies = rows[:limit] return PolicyPage( items=tuple(policy_read(policy, values) for policy in policies), @@ -210,6 +224,11 @@ def _utc_midnight(value: str) -> datetime: def parse_legacy_period(value: str) -> tuple[datetime, datetime]: + """Return the UTC-midnight start and end dates of a v1 period key. + + Accept a ``start.end`` ISO date range or a PolicyEngine period such as + ``2024``. Raise LegacyPolicyTranslationError for anything else. + """ if not value or value != value.strip(): raise LegacyPolicyTranslationError("legacy period must be non-empty") if "." in value: @@ -242,6 +261,12 @@ def translate_legacy_policy( parameters: Mapping[str, Parameter], country_package_versions: Mapping[str, str], ) -> ResolvedPolicyCreationInput: + """Translate a v1 policy into policy content bound to the running catalog. + + Raise LegacyPolicyTranslationError if the v1 policy was saved under another + country package version, names a parameter the catalog lacks, or has + invalid or conflicting periods. + """ expected_api_version = country_package_versions.get(snapshot.country_id) if expected_api_version is None or snapshot.api_version != expected_api_version: raise LegacyPolicyTranslationError( diff --git a/policyengine_api/services/v2/policies/types.py b/policyengine_api/services/v2/policies/types.py index ef1f2823a..fcf1647c3 100644 --- a/policyengine_api/services/v2/policies/types.py +++ b/policyengine_api/services/v2/policies/types.py @@ -45,6 +45,7 @@ class PolicyParameterValueInput(StrictPolicyInput): @model_validator(mode="after") def validate_period(self) -> PolicyParameterValueInput: + """Reject a value whose end date precedes its start date.""" if self.end_date is not None and self.end_date < self.start_date: raise ValueError("end_date must not precede start_date") return self @@ -66,6 +67,7 @@ def reject_duplicate_effective_values( cls, values: list[PolicyParameterValueInput], ) -> list[PolicyParameterValueInput]: + """Reject two values for the same parameter with the same start date.""" identities: set[tuple[UUID, datetime]] = set() for value in values: identity = (value.parameter_id, value.start_date) @@ -103,6 +105,7 @@ class LegacyPolicySnapshot(StrictPolicyInput): @field_validator("policy_json") @classmethod def require_parameter_mapping(cls, value: object) -> object: + """Reject a v1 policy_json that is not an object keyed by parameter path.""" if type(value) is not dict: raise ValueError("legacy policy_json must be a parameter-path object") return value @@ -136,6 +139,8 @@ class LegacyPolicyPersistenceResult: @dataclass(frozen=True) class PolicyParameterValueRead: + """One stored parameter value of a policy, with the parameter's name.""" + id: UUID parameter_id: UUID parameter_name: str @@ -146,6 +151,8 @@ class PolicyParameterValueRead: @dataclass(frozen=True) class PolicyRead: + """A stored v2 policy with its parameter values, as returned to routes.""" + id: UUID country_id: str tax_benefit_model_id: UUID @@ -157,6 +164,8 @@ class PolicyRead: @dataclass(frozen=True) class PolicyPage: + """One page of policies in creation order; ``has_more`` says if more follow.""" + items: tuple[PolicyRead, ...] offset: int limit: int diff --git a/policyengine_api/services/v2/reports/types.py b/policyengine_api/services/v2/reports/types.py index 8cb7041f1..cd6d8d2b1 100644 --- a/policyengine_api/services/v2/reports/types.py +++ b/policyengine_api/services/v2/reports/types.py @@ -22,6 +22,8 @@ class ReportAggregate(StrEnum): + """An aggregate result section a report can request.""" + BUDGET = "budget" POVERTY = "poverty" INEQUALITY = "inequality" @@ -32,6 +34,12 @@ class ReportAggregate(StrEnum): class ReportExecutionInput(StrictContractModel): + """The complete input for one report run. + + It pairs a baseline and a reform simulation that differ only in policy, and + lists the aggregates to compute from them. + """ + contract_version: ContractVersion = 1 evaluation_id: UUID baseline: SimulationExecutionInput @@ -43,6 +51,13 @@ class ReportExecutionInput(StrictContractModel): @model_validator(mode="after") def require_aligned_simulations(self) -> ReportExecutionInput: + """Reject a report whose two simulations cannot be compared. + + Baseline and reform must carry their own roles and this report's + evaluation_id, have distinct execution ids, and share population, year, + geography, options, requested output and bundle. Requested aggregates + must be unique. + """ if self.baseline.role is not SimulationRole.BASELINE: raise ValueError("baseline input must use the baseline role") if self.reform.role is not SimulationRole.REFORM: @@ -72,6 +87,12 @@ def require_aligned_simulations(self) -> ReportExecutionInput: class AggregateReportArtifactDescriptor(StrictContractModel): + """The receipt for one stored aggregate report. + + It names the artifact and the bundle, and records the digests of the + baseline and reform simulation artifacts it was computed from. + """ + contract_version: ContractVersion = 1 evaluation_id: UUID artifact: ArtifactReference @@ -100,6 +121,7 @@ def require_unique_aggregates( cls, value: tuple[ReportAggregate, ...], ) -> tuple[ReportAggregate, ...]: + """Reject a payload that lists any aggregate twice.""" if len(value) != len(set(value)): raise ValueError("requested aggregates must be unique") return value diff --git a/policyengine_api/services/v2/simulations/types.py b/policyengine_api/services/v2/simulations/types.py index d1e466b13..7a19b9793 100644 --- a/policyengine_api/services/v2/simulations/types.py +++ b/policyengine_api/services/v2/simulations/types.py @@ -46,12 +46,16 @@ class StrictContractModel(BaseModel): class SimulationRole(StrEnum): + """The part a simulation plays: one side of a comparison, or a standalone run.""" + BASELINE = "baseline" REFORM = "reform" STANDALONE = "standalone" class ArtifactMediaType(StrEnum): + """Media type of a worker output: JSON for reports, Parquet for simulations.""" + JSON = "application/json" PARQUET = "application/vnd.apache.parquet" @@ -90,12 +94,16 @@ class SimulationParquetPayloadContract(StrictContractModel): class DatasetArtifactMediaType(StrEnum): + """Media type of an input dataset artifact.""" + HDF5 = "application/x-hdf5" JSON = "application/json" PARQUET = "application/vnd.apache.parquet" class DatasetProvenance(StrictContractModel): + """The input dataset a simulation runs on, with its revision and data package.""" + identity: ContractText uri: Annotated[str, Field(min_length=1, max_length=2048)] artifact_revision: ContractText @@ -104,6 +112,12 @@ class DatasetProvenance(StrictContractModel): class BundleProvenance(StrictContractModel): + """The PolicyEngine.py bundle a simulation runs against. + + It pins the package versions, the input dataset and the bundle manifest + digest, so a result can be traced to an exact release. + """ + policyengine_version: ContractText country_package_name: ContractText country_package_version: ContractText @@ -112,6 +126,8 @@ class BundleProvenance(StrictContractModel): class ArtifactReference(StrictContractModel): + """A private gs:// object, addressed by its SHA-256 digest and size.""" + uri: StorageUri media_type: ArtifactMediaType content_sha256: Sha256Digest @@ -119,6 +135,11 @@ class ArtifactReference(StrictContractModel): class DatasetArtifactReference(StrictContractModel): + """An input dataset object in GCS or Hugging Face, addressed by its SHA-256 digest. + + size_bytes is optional because a dataset's size may not be known. + """ + uri: DatasetArtifactUri media_type: DatasetArtifactMediaType content_sha256: Sha256Digest @@ -126,11 +147,15 @@ class DatasetArtifactReference(StrictContractModel): class DatasetPopulationInput(StrictContractModel): + """A population supplied as an input dataset artifact.""" + kind: Literal["dataset"] = "dataset" artifact: DatasetArtifactReference class HouseholdPopulationInput(StrictContractModel): + """A population supplied as one household document.""" + kind: Literal["household"] = "household" document: dict[str, JsonValue] @@ -142,6 +167,12 @@ class HouseholdPopulationInput(StrictContractModel): class GeographySelection(StrictContractModel): + """The country and region to simulate. + + filter_field and filter_value, set together, restrict the population to the + region; filter_strategy, if set, names how the filter is applied. + """ + country: CountryId region: ContractText filter_field: ContractText | None = None @@ -150,6 +181,7 @@ class GeographySelection(StrictContractModel): @model_validator(mode="after") def require_complete_filter(self) -> GeographySelection: + """Require filter_field with filter_value, and both for filter_strategy.""" if (self.filter_field is None) != (self.filter_value is None): raise ValueError("filter_field and filter_value must be supplied together") if self.filter_strategy is not None and self.filter_field is None: @@ -158,18 +190,27 @@ def require_complete_filter(self) -> GeographySelection: class RequestedSimulationOutput(StrictContractModel): + """The variables the worker writes to a simulation's output, each named once.""" + schema_version: ContractVersion = 1 variables: Annotated[tuple[ContractText, ...], Field(min_length=1)] @field_validator("variables") @classmethod def require_unique_variables(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Reject a variable list that names any variable twice.""" if len(value) != len(set(value)): raise ValueError("requested variables must be unique") return value class SimulationExecutionInput(StrictContractModel): + """The complete input for one simulation in a comparison run. + + It fixes the policy, population, year, geography, options, requested + outputs and the bundle to run against. + """ + contract_version: ContractVersion = 1 evaluation_id: UUID simulation_execution_id: UUID @@ -184,6 +225,8 @@ class SimulationExecutionInput(StrictContractModel): class RowIdentity(StrictContractModel): + """The identifier columns, row count and digest that fix the rows of an output.""" + schema_version: ContractVersion = 1 identifier_columns: Annotated[tuple[ContractText, ...], Field(min_length=1)] row_count: Annotated[int, Field(ge=0)] @@ -192,12 +235,19 @@ class RowIdentity(StrictContractModel): @field_validator("identifier_columns") @classmethod def require_unique_columns(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Reject identifier columns that name any column twice.""" if len(value) != len(set(value)): raise ValueError("row identifier columns must be unique") return value class SimulationArtifactDescriptor(StrictContractModel): + """The receipt for one stored simulation output. + + It names the artifact, its row identity and bundle, and any detached SPM + calculation receipt needed to reproduce an aggregate from it. + """ + contract_version: ContractVersion = 1 evaluation_id: UUID simulation_execution_id: UUID diff --git a/policyengine_api/services/v2/user_households/database_connectors/creates.py b/policyengine_api/services/v2/user_households/database_connectors/creates.py index eb871996f..4b80c3a7f 100644 --- a/policyengine_api/services/v2/user_households/database_connectors/creates.py +++ b/policyengine_api/services/v2/user_households/database_connectors/creates.py @@ -12,6 +12,7 @@ def create_user_household( session: Session, association_input: UserHouseholdCreationInput, ) -> UserHouseholdAssociation: + """Insert an association and return it with its database-generated fields loaded.""" association = UserHouseholdAssociation(**association_input.model_dump()) session.add(association) session.flush() diff --git a/policyengine_api/services/v2/user_households/database_connectors/deletes.py b/policyengine_api/services/v2/user_households/database_connectors/deletes.py index 800a3e645..b18a2ccfc 100644 --- a/policyengine_api/services/v2/user_households/database_connectors/deletes.py +++ b/policyengine_api/services/v2/user_households/database_connectors/deletes.py @@ -9,5 +9,6 @@ def delete_user_household( session: Session, association: UserHouseholdAssociation, ) -> None: + """Delete an association row; the household it points to is kept.""" session.delete(association) session.flush() diff --git a/policyengine_api/services/v2/user_households/database_connectors/reads.py b/policyengine_api/services/v2/user_households/database_connectors/reads.py index bcab40a73..b6b4df7fd 100644 --- a/policyengine_api/services/v2/user_households/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_households/database_connectors/reads.py @@ -13,6 +13,7 @@ def read_user(session: Session, user_id: UUID) -> User | None: + """Return the v2 user with this id, or None.""" return session.get(User, user_id) @@ -20,6 +21,7 @@ def read_household_for_association( session: Session, household_id: UUID, ) -> Household | None: + """Return the household with this id in any country, or None.""" return session.exec( select(Household).where(Household.id == household_id) ).one_or_none() @@ -31,6 +33,7 @@ def read_user_household_row( country_id: str, association_id: UUID, ) -> UserHouseholdAssociation | None: + """Return the association with this id in this country, or None.""" return session.exec( select(UserHouseholdAssociation).where( UserHouseholdAssociation.id == association_id, @@ -48,6 +51,11 @@ def read_user_household_rows( offset: int, limit: int, ) -> list[UserHouseholdAssociation]: + """Return up to ``limit + 1`` of a user's associations in creation order. + + With ``household_id``, only associations to that household are returned. The + extra row lets the caller report whether another page exists. + """ statement = select(UserHouseholdAssociation).where( UserHouseholdAssociation.country_id == country_id, UserHouseholdAssociation.user_id == user_id, diff --git a/policyengine_api/services/v2/user_households/database_connectors/updates.py b/policyengine_api/services/v2/user_households/database_connectors/updates.py index 2e7c6a030..59fc2f761 100644 --- a/policyengine_api/services/v2/user_households/database_connectors/updates.py +++ b/policyengine_api/services/v2/user_households/database_connectors/updates.py @@ -14,6 +14,7 @@ def update_user_household( association: UserHouseholdAssociation, association_input: UserHouseholdUpdateInput, ) -> UserHouseholdAssociation: + """Apply only the supplied fields, bump ``updated_at`` and return the row.""" for field_name, value in association_input.model_dump(exclude_unset=True).items(): setattr(association, field_name, value) association.updated_at = utc_now() diff --git a/policyengine_api/services/v2/user_households/database_session.py b/policyengine_api/services/v2/user_households/database_session.py index ffa531976..2f7fae41d 100644 --- a/policyengine_api/services/v2/user_households/database_session.py +++ b/policyengine_api/services/v2/user_households/database_session.py @@ -13,15 +13,20 @@ class UserHouseholdDatabaseSession: + """Session lifetimes for user-household services: reads and transactions.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/user_households/services.py b/policyengine_api/services/v2/user_households/services.py index a26a9bf5c..5cb2d7fd9 100644 --- a/policyengine_api/services/v2/user_households/services.py +++ b/policyengine_api/services/v2/user_households/services.py @@ -51,6 +51,7 @@ def read_complete_user_household( country_id: str, association_id: UUID, ) -> UserHouseholdRead: + """Return this country's association, or raise UserHouseholdNotFoundError.""" association = require_user_household( read_user_household_row( session, @@ -71,6 +72,7 @@ def read_user_household_page( offset: int, limit: int, ) -> UserHouseholdPage: + """Return a page of a user's associations, optionally for one household.""" rows = read_user_household_rows( session, country_id=country_id, @@ -86,12 +88,19 @@ class V2UserHouseholdService: """Sequence association operations through explicit database sessions.""" def __init__(self, database_session: UserHouseholdDatabaseSession) -> None: + """Open every session through ``database_session``.""" self._database_session = database_session def create_user_household( self, association_input: UserHouseholdCreationInput, ) -> UserHouseholdRead: + """Save a link from an existing v2 user to an existing household. + + The user id is not authenticated. Raise AssociationUserNotFoundError, + AssociationHouseholdNotFoundError or AssociationCountryConflictError when the + user or household is missing or the household is from another country. + """ with self._database_session.transaction() as session: household = read_household_for_association( session, @@ -112,6 +121,7 @@ def get_user_household( country_id: str, association_id: UUID, ) -> UserHouseholdRead: + """Return this country's association, or raise UserHouseholdNotFoundError.""" with self._database_session.read() as session: return read_complete_user_household( session, @@ -128,6 +138,7 @@ def list_user_households( offset: int = 0, limit: int = 100, ) -> UserHouseholdPage: + """Return a page of a user's associations, optionally for one household.""" with self._database_session.read() as session: return read_user_household_page( session, @@ -145,6 +156,10 @@ def patch_user_household( association_id: UUID, association_input: UserHouseholdUpdateInput, ) -> UserHouseholdRead: + """Update one association's supplied fields, including its household. + + A replacement household must exist and belong to the association's country. + """ with self._database_session.transaction() as session: association = require_user_household( read_user_household_row( @@ -175,6 +190,7 @@ def delete_user_household( country_id: str, association_id: UUID, ) -> None: + """Delete this country's association, or raise UserHouseholdNotFoundError.""" with self._database_session.transaction() as session: association = require_user_household( read_user_household_row( diff --git a/policyengine_api/services/v2/user_households/transformations.py b/policyengine_api/services/v2/user_households/transformations.py index 1c0203573..98698db43 100644 --- a/policyengine_api/services/v2/user_households/transformations.py +++ b/policyengine_api/services/v2/user_households/transformations.py @@ -10,6 +10,7 @@ def user_household_read( association: UserHouseholdAssociation, ) -> UserHouseholdRead: + """Convert a stored association row to the service's read model.""" return UserHouseholdRead( id=association.id, country_id=association.country_id, @@ -28,6 +29,10 @@ def user_household_page( offset: int, limit: int, ) -> UserHouseholdPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ return UserHouseholdPage( items=tuple(user_household_read(row) for row in rows[:limit]), offset=offset, diff --git a/policyengine_api/services/v2/user_households/types.py b/policyengine_api/services/v2/user_households/types.py index 7ce3c1fb2..d968eb01c 100644 --- a/policyengine_api/services/v2/user_households/types.py +++ b/policyengine_api/services/v2/user_households/types.py @@ -17,10 +17,14 @@ class StrictUserHouseholdInput(BaseModel): + """Base for association inputs that rejects undeclared fields and is immutable.""" + model_config = ConfigDict(extra="forbid", frozen=True) class UserHouseholdCreationInput(StrictUserHouseholdInput): + """A new saved link between an existing v2 user and an existing household.""" + country_id: CountryId user_id: UserId household_id: ResourceId @@ -29,12 +33,19 @@ class UserHouseholdCreationInput(StrictUserHouseholdInput): class UserHouseholdUpdateInput(StrictUserHouseholdInput): + """A partial update of an association; only supplied fields are changed. + + At least one field must be supplied, and household_id may be changed but not + cleared. + """ + household_id: ResourceId | None = None name: AssociationName | None = None description: str | None = None @model_validator(mode="after") def require_valid_supplied_field(self) -> UserHouseholdUpdateInput: + """Require a supplied field, and reject an explicit null household_id.""" mutable_fields = {"household_id", "name", "description"} supplied = self.model_fields_set.intersection(mutable_fields) if not supplied: @@ -48,6 +59,8 @@ def require_valid_supplied_field(self) -> UserHouseholdUpdateInput: @dataclass(frozen=True) class UserHouseholdRead: + """A stored user-household association as the service returns it to routes.""" + id: UUID country_id: str user_id: UUID @@ -60,6 +73,8 @@ class UserHouseholdRead: @dataclass(frozen=True) class UserHouseholdPage: + """One page of associations in creation order; ``has_more`` says if more follow.""" + items: tuple[UserHouseholdRead, ...] offset: int limit: int diff --git a/policyengine_api/services/v2/user_households/validators.py b/policyengine_api/services/v2/user_households/validators.py index 4a27aa531..4a1a5afe7 100644 --- a/policyengine_api/services/v2/user_households/validators.py +++ b/policyengine_api/services/v2/user_households/validators.py @@ -38,6 +38,7 @@ def require_user_household( *, association_id: UUID, ) -> UserHouseholdAssociation: + """Return the association, or raise UserHouseholdNotFoundError if it is None.""" if association is None: raise UserHouseholdNotFoundError( f"user-household association {association_id} was not found" @@ -51,6 +52,7 @@ def validate_association_creation( user: User | None, household: Household | None, ) -> None: + """Check that the referenced user exists and the household fits the association.""" if user is None: raise AssociationUserNotFoundError( f"user {association_input.user_id} was not found" @@ -68,6 +70,10 @@ def validate_association_household( household_id: UUID, household: Household | None, ) -> None: + """Check that a referenced household exists and is in the association's country. + + Raise AssociationHouseholdNotFoundError or AssociationCountryConflictError. + """ if household is None: raise AssociationHouseholdNotFoundError( f"household {household_id} was not found" diff --git a/policyengine_api/services/v2/user_policies/database_connectors/creates.py b/policyengine_api/services/v2/user_policies/database_connectors/creates.py index 822abafd3..4ed880968 100644 --- a/policyengine_api/services/v2/user_policies/database_connectors/creates.py +++ b/policyengine_api/services/v2/user_policies/database_connectors/creates.py @@ -23,6 +23,7 @@ def create_user_policy( session: Session, association_input: UserPolicyCreationInput ) -> UserPolicy: + """Insert an association and return it with its database-generated fields loaded.""" association = UserPolicy(**association_input.model_dump()) session.add(association) session.flush() @@ -31,6 +32,7 @@ def create_user_policy( def create_transition_user(session: Session, *, primary_country: str) -> User: + """Insert a v2 user to stand in for a v1 user that has no v2 user yet.""" user = User(primary_country=primary_country) session.add(user) session.flush() @@ -40,6 +42,11 @@ def create_transition_user(session: Session, *, primary_country: str) -> User: def create_legacy_user_mapping( session: Session, *, legacy_user_id: str, user_id: UUID ) -> UUID | None: + """Map a v1 user id to a v2 user. + + Return the v2 user id, or None when the v1 id is already mapped (for example, + a concurrent copy won the race). + """ return session.execute( insert(LegacyUserMapping) .values(legacy_user_id=legacy_user_id, user_id=user_id) @@ -58,6 +65,11 @@ def create_legacy_user_policy_mapping( fingerprint_version: int, fingerprint: str, ) -> UUID | None: + """Map a v1 saved policy to its v2 association at ``source_revision``. + + Return the new mapping id, or None when this v1 saved policy is already + mapped in this country. + """ return session.execute( insert(LegacyUserPolicyMapping) .values( diff --git a/policyengine_api/services/v2/user_policies/database_connectors/deletes.py b/policyengine_api/services/v2/user_policies/database_connectors/deletes.py index 7b62bb24b..d9f2bc6be 100644 --- a/policyengine_api/services/v2/user_policies/database_connectors/deletes.py +++ b/policyengine_api/services/v2/user_policies/database_connectors/deletes.py @@ -11,10 +11,12 @@ def delete_user_policy(session: Session, association: UserPolicy) -> None: + """Delete an association row; the policy it points to is kept.""" session.delete(association) session.flush() def delete_transition_user(session: Session, user: User) -> None: + """Delete a transition user that lost a race to map its v1 user id.""" session.delete(user) session.flush() diff --git a/policyengine_api/services/v2/user_policies/database_connectors/reads.py b/policyengine_api/services/v2/user_policies/database_connectors/reads.py index 0ac9acaa5..e2b01a6fa 100644 --- a/policyengine_api/services/v2/user_policies/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_policies/database_connectors/reads.py @@ -19,16 +19,19 @@ def read_user(session: Session, user_id: UUID) -> User | None: + """Return the v2 user with this id, or None.""" return session.get(User, user_id) def read_policy_for_association(session: Session, policy_id: UUID) -> Policy | None: + """Return the policy with this id in any country, or None.""" return session.exec(select(Policy).where(Policy.id == policy_id)).one_or_none() def read_legacy_user_mapping( session: Session, legacy_user_id: str, *, lock: bool ) -> LegacyUserMapping | None: + """Return the mapping for a v1 user id, or None; ``lock`` selects it FOR UPDATE.""" statement = select(LegacyUserMapping).where( LegacyUserMapping.legacy_user_id == legacy_user_id ) @@ -44,6 +47,11 @@ def read_legacy_user_policy_mapping( legacy_user_policy_id: int, lock: bool, ) -> LegacyUserPolicyMapping | None: + """Return the mapping for a v1 saved policy in this country, or None. + + With ``lock``, the row is selected FOR UPDATE so concurrent revisions of the + same v1 saved policy serialize on it. + """ statement = select(LegacyUserPolicyMapping).where( LegacyUserPolicyMapping.country_id == country_id, LegacyUserPolicyMapping.legacy_user_policy_id == legacy_user_policy_id, @@ -56,6 +64,7 @@ def read_legacy_user_policy_mapping( def read_mapped_user_policy( session: Session, mapping: LegacyUserPolicyMapping ) -> UserPolicy | None: + """Return the association a mapping points to, in the mapping's country, or None.""" return session.exec( select(UserPolicy).where( UserPolicy.id == mapping.user_policy_id, @@ -67,6 +76,7 @@ def read_mapped_user_policy( def read_user_policy_row( session: Session, *, country_id: str, association_id: UUID ) -> UserPolicy | None: + """Return the association with this id in this country, or None.""" return session.exec( select(UserPolicy).where( UserPolicy.id == association_id, @@ -84,6 +94,11 @@ def read_user_policy_rows( offset: int, limit: int, ) -> list[UserPolicy]: + """Return up to ``limit + 1`` of a user's associations in creation order. + + With ``policy_id``, only associations to that policy are returned. The extra + row lets the caller report whether another page exists. + """ statement = select(UserPolicy).where( UserPolicy.country_id == country_id, UserPolicy.user_id == user_id, diff --git a/policyengine_api/services/v2/user_policies/database_connectors/updates.py b/policyengine_api/services/v2/user_policies/database_connectors/updates.py index 0422b1317..15c5a8216 100644 --- a/policyengine_api/services/v2/user_policies/database_connectors/updates.py +++ b/policyengine_api/services/v2/user_policies/database_connectors/updates.py @@ -18,6 +18,7 @@ def update_user_policy( association: UserPolicy, association_input: UserPolicyUpdateInput, ) -> UserPolicy: + """Apply only the supplied fields, bump ``updated_at`` and return the row.""" for field_name, value in association_input.model_dump(exclude_unset=True).items(): setattr(association, field_name, value) association.updated_at = utc_now() @@ -37,6 +38,10 @@ def update_legacy_user_policy_state( fingerprint: str, source_revision: int, ) -> None: + """Record a newly applied v1 revision on its mapping. + + With ``update_name``, also set the association's name to the v1 reform label. + """ if update_name: association.name = reform_label association.updated_at = utc_now() diff --git a/policyengine_api/services/v2/user_policies/database_session.py b/policyengine_api/services/v2/user_policies/database_session.py index 78d084ac9..8ec71003f 100644 --- a/policyengine_api/services/v2/user_policies/database_session.py +++ b/policyengine_api/services/v2/user_policies/database_session.py @@ -16,14 +16,17 @@ class UserPolicyDatabaseSession: """Open read sessions and transactions for user-policy services.""" def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/user_policies/services.py b/policyengine_api/services/v2/user_policies/services.py index 627ea2898..9d7e3fc27 100644 --- a/policyengine_api/services/v2/user_policies/services.py +++ b/policyengine_api/services/v2/user_policies/services.py @@ -72,6 +72,7 @@ def read_complete_user_policy( session: Session, *, country_id: str, association_id: UUID ) -> UserPolicyRead: + """Return this country's association, or raise UserPolicyNotFoundError.""" association = require_user_policy( read_user_policy_row( session, @@ -92,6 +93,7 @@ def read_user_policy_page( offset: int, limit: int, ) -> UserPolicyPage: + """Return a page of a user's associations, optionally for one policy.""" rows = read_user_policy_rows( session, country_id=country_id, @@ -109,6 +111,11 @@ def resolve_legacy_user_id( legacy_user_id: str, primary_country: str, ) -> UUID: + """Return the v2 user id for a v1 user id, creating a transition user if needed. + + If a concurrent transaction maps the v1 id first, the new user is deleted and + the winner's user is returned. + """ existing = read_legacy_user_mapping(session, legacy_user_id, lock=True) if existing is not None: return validate_existing_legacy_user_mapping( @@ -147,6 +154,12 @@ def apply_existing_legacy_user_policy_mapping( changed_fields: frozenset[str], source_revision: int, ) -> LegacyUserPolicyPersistenceResult: + """Apply one v1 revision to an already-mapped association. + + Stale and replayed revisions change nothing. The next revision renames the + association only if the v1 reform label changed. Raise + LegacyUserPolicyIntegrityError on conflicting or skipped revisions. + """ action, association = validate_existing_legacy_user_policy_mapping( mapping, association=read_mapped_user_policy(session, mapping), @@ -203,6 +216,11 @@ def persist_legacy_user_policy_mapping( source_revision: int, changed_fields: frozenset[str], ) -> LegacyUserPolicyPersistenceResult: + """Create the association and mapping for a v1 saved policy, or update them. + + If a concurrent transaction creates the mapping first, the new association is + deleted and the revision is applied to the winner's association. + """ existing = read_legacy_user_policy_mapping( session, country_id=country_id, @@ -273,6 +291,11 @@ def mirror_legacy_user_policy_in_session( source_revision: int, changed_fields: frozenset[str] = frozenset(), ) -> LegacyUserPolicyPersistenceResult: + """Copy one revision of a v1 saved policy into v2. + + Copy its reform policy, resolve the v1 user to a v2 user, then create or + update the association. Only the v1 reform label is copied, as the name. + """ validate_legacy_user_policy_input( snapshot, reform_snapshot, @@ -305,11 +328,18 @@ class V2UserPolicyService: """Sequence association operations through explicit database sessions.""" def __init__(self, database_session: UserPolicyDatabaseSession) -> None: + """Open every session through ``database_session``.""" self._database_session = database_session def create_user_policy( self, association_input: UserPolicyCreationInput ) -> UserPolicyRead: + """Save a link from an existing v2 user to an existing policy. + + The user id is not authenticated. Raise AssociationUserNotFoundError, + AssociationPolicyNotFoundError or AssociationCountryConflictError when the + user or policy is missing or the policy is from another country. + """ with self._database_session.transaction() as session: policy = read_policy_for_association(session, association_input.policy_id) validate_association_creation( @@ -322,6 +352,7 @@ def create_user_policy( def get_user_policy( self, *, country_id: str, association_id: UUID ) -> UserPolicyRead: + """Return this country's association, or raise UserPolicyNotFoundError.""" with self._database_session.read() as session: return read_complete_user_policy( session, @@ -338,6 +369,7 @@ def list_user_policies( offset: int = 0, limit: int = 100, ) -> UserPolicyPage: + """Return a page of a user's associations, optionally for one policy.""" with self._database_session.read() as session: return read_user_policy_page( session, @@ -355,6 +387,7 @@ def patch_user_policy( association_id: UUID, association_input: UserPolicyUpdateInput, ) -> UserPolicyRead: + """Update the supplied name or description; the linked policy cannot change.""" with self._database_session.transaction() as session: association = require_user_policy( read_user_policy_row( @@ -369,6 +402,7 @@ def patch_user_policy( ) def delete_user_policy(self, *, country_id: str, association_id: UUID) -> None: + """Delete this country's association, or raise UserPolicyNotFoundError.""" with self._database_session.transaction() as session: association = require_user_policy( read_user_policy_row( @@ -388,6 +422,7 @@ def mirror_legacy_user_policy( source_revision: int, changed_fields: frozenset[str], ) -> LegacyUserPolicyPersistenceResult: + """Copy one v1 saved-policy revision into v2 in its own transaction.""" with self._database_session.transaction() as session: return mirror_legacy_user_policy_in_session( session, diff --git a/policyengine_api/services/v2/user_policies/transformations.py b/policyengine_api/services/v2/user_policies/transformations.py index 32f846e21..238a66c83 100644 --- a/policyengine_api/services/v2/user_policies/transformations.py +++ b/policyengine_api/services/v2/user_policies/transformations.py @@ -22,6 +22,7 @@ def user_policy_read(association: UserPolicy) -> UserPolicyRead: + """Convert a stored association row to the service's read model.""" return UserPolicyRead( id=association.id, country_id=association.country_id, @@ -37,6 +38,10 @@ def user_policy_read(association: UserPolicy) -> UserPolicyRead: def user_policy_page( rows: list[UserPolicy], *, offset: int, limit: int ) -> UserPolicyPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ return UserPolicyPage( items=tuple(user_policy_read(row) for row in rows[:limit]), offset=offset, @@ -83,4 +88,9 @@ def legacy_name_requires_update( reform_label: str | None, changed_fields: frozenset[str], ) -> bool: + """Return whether a v1 revision should rename the association. + + Only a changed v1 reform label that differs from the current name renames it, + so a v2-side rename survives v1 revisions that leave the label alone. + """ return "reform_label" in changed_fields and association.name != reform_label diff --git a/policyengine_api/services/v2/user_policies/types.py b/policyengine_api/services/v2/user_policies/types.py index 2a1a62ded..0e80ac0fb 100644 --- a/policyengine_api/services/v2/user_policies/types.py +++ b/policyengine_api/services/v2/user_policies/types.py @@ -27,6 +27,8 @@ class StrictAssociationInput(BaseModel): class UserPolicyCreationInput(StrictAssociationInput): + """A new saved link between an existing v2 user and an existing policy.""" + country_id: CountryId user_id: UserId policy_id: ResourceId @@ -35,11 +37,17 @@ class UserPolicyCreationInput(StrictAssociationInput): class UserPolicyUpdateInput(StrictAssociationInput): + """A partial update of an association's name or description. + + At least one of them must be supplied; the linked policy cannot change. + """ + name: Annotated[str, StringConstraints(max_length=255)] | None = None description: str | None = None @model_validator(mode="after") def require_supplied_field(self) -> UserPolicyUpdateInput: + """Require at least one of name or description to be supplied.""" if not self.model_fields_set.intersection({"name", "description"}): raise ValueError("At least one of name or description must be supplied") return self @@ -67,6 +75,12 @@ class LegacyUserPolicySnapshot(StrictAssociationInput): class LegacyUserPolicyMappingAction(StrEnum): + """What to do with a v1 saved-policy revision whose row is already mapped. + + STALE is older than the applied revision and is ignored, REPLAY repeats the + applied revision, and UPDATE is the next revision, which is applied. + """ + STALE = "stale" REPLAY = "replay" UPDATE = "update" @@ -74,6 +88,8 @@ class LegacyUserPolicyMappingAction(StrEnum): @dataclass(frozen=True) class UserPolicyRead: + """A stored user-policy association as the service returns it to routes.""" + id: UUID country_id: str user_id: UUID @@ -86,6 +102,8 @@ class UserPolicyRead: @dataclass(frozen=True) class UserPolicyPage: + """One page of associations in creation order; ``has_more`` says if more follow.""" + items: tuple[UserPolicyRead, ...] offset: int limit: int @@ -94,6 +112,12 @@ class UserPolicyPage: @dataclass(frozen=True) class LegacyUserPolicyPersistenceResult: + """Outcome of copying one v1 saved-policy revision into v2. + + Names the v2 association and policy, and whether this call created the + association, renamed it, or created the v1-to-v2 mapping. + """ + association_id: UUID policy_id: UUID association_created: bool diff --git a/policyengine_api/services/v2/user_policies/validators.py b/policyengine_api/services/v2/user_policies/validators.py index 5a5916c09..5e33b74e5 100644 --- a/policyengine_api/services/v2/user_policies/validators.py +++ b/policyengine_api/services/v2/user_policies/validators.py @@ -48,6 +48,7 @@ class LegacyUserPolicyIntegrityError(RuntimeError): def require_user_policy( association: UserPolicy | None, *, association_id: UUID ) -> UserPolicy: + """Return the association, or raise UserPolicyNotFoundError if it is None.""" if association is None: raise UserPolicyNotFoundError( f"user-policy association {association_id} was not found" @@ -61,6 +62,11 @@ def validate_association_creation( user: User | None, policy: Policy | None, ) -> None: + """Check that the referenced user and policy exist and share the country. + + Raise AssociationUserNotFoundError, AssociationPolicyNotFoundError or + AssociationCountryConflictError. + """ if user is None: raise AssociationUserNotFoundError( f"user {association_input.user_id} was not found" @@ -78,6 +84,7 @@ def validate_association_creation( def validate_existing_legacy_user_mapping( mapping: LegacyUserMapping, *, user: User | None ) -> UUID: + """Return the mapped v2 user id, or raise if that user no longer exists.""" if user is None: raise LegacyUserPolicyIntegrityError( "legacy user mapping has no referenced v2 user" @@ -88,6 +95,7 @@ def validate_existing_legacy_user_mapping( def validate_legacy_user_mapping_conflict( mapping: LegacyUserMapping | None, *, user: User | None ) -> UUID: + """Return the user id of the mapping that won the race, or raise if unusable.""" if mapping is None or user is None: raise LegacyUserPolicyIntegrityError( "legacy user mapping conflict did not resolve to a v2 user" @@ -101,6 +109,7 @@ def validate_legacy_user_policy_input( *, source_revision: int, ) -> None: + """Require a positive revision and a saved policy that uses the reform snapshot.""" if source_revision <= 0: raise LegacyUserPolicyIntegrityError( "legacy user-policy source revision must be positive" @@ -125,6 +134,13 @@ def validate_existing_legacy_user_policy_mapping( source_revision: int, fingerprint_version: int, ) -> tuple[LegacyUserPolicyMappingAction, UserPolicy]: + """Decide how to apply a v1 revision to an existing mapping. + + Return the action with the mapped association. Raise + LegacyUserPolicyIntegrityError if the association is missing or its user, + policy or country changed, if the fingerprint version differs, if a replay + has a different fingerprint, or if a revision was skipped. + """ if association is None: raise LegacyUserPolicyIntegrityError( "legacy user-policy mapping has no association" From fefcddbc84eadab22bf1091b5fe7ec29442a9368 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:23:52 +0400 Subject: [PATCH 07/12] Document the FastAPI routes, dependencies and response models Add docstrings to the native route dependencies (one contract line per Protocol method), the health, metadata and specification routers, the v2 metadata route builders and their shared read_resource helper, and the v2 request, document and response models and their methods. Pydantic publishes a model's class docstring as its JSON-schema description, so the response and document models in the OpenAPI document now carry descriptions written for API consumers: what each object represents, that households and policies are immutable and deduplicated, and how to page through results. Route functions, whose docstrings FastAPI would publish as operation descriptions, are unchanged, and so are the documented paths and operations. Co-Authored-By: Claude Opus 5.5 --- .../fastapi_routes/dependencies.py | 194 ++++++++++++++---- policyengine_api/fastapi_routes/health.py | 4 + policyengine_api/fastapi_routes/metadata.py | 2 + policyengine_api/fastapi_routes/responses.py | 1 + .../v2/households/document_models.py | 29 +++ .../v2/households/request_models.py | 2 + .../v2/households/response_models.py | 24 ++- .../fastapi_routes/v2/metadata/common.py | 6 + .../v2/metadata/geography_routes.py | 5 + .../v2/metadata/model_routes.py | 5 + .../v2/metadata/parameter_routes.py | 5 + .../v2/metadata/response_models.py | 42 ++-- .../v2/policies/response_models.py | 28 ++- .../v2/user_households/response_models.py | 24 ++- .../v2/user_policies/response_models.py | 21 +- 15 files changed, 329 insertions(+), 63 deletions(-) diff --git a/policyengine_api/fastapi_routes/dependencies.py b/policyengine_api/fastapi_routes/dependencies.py index d13e2649d..028d1e6df 100644 --- a/policyengine_api/fastapi_routes/dependencies.py +++ b/policyengine_api/fastapi_routes/dependencies.py @@ -56,13 +56,24 @@ class MetadataReader(Protocol): """Read the already-loaded metadata document for a country.""" - def get_metadata(self, country_id: str) -> JSONObject: ... + def get_metadata(self, country_id: str) -> JSONObject: + """Return the loaded v1 metadata document for a supported country.""" + ... class V2MetadataResourceReader(Protocol): - """Own the database session used by one v2 metadata resource request.""" + """Own the database session used by one v2 metadata resource request. - def close(self) -> None: ... + Every read selects one country catalog: the one published for + ``policyengine_version``, or the running package's catalog when it is None. + Absent resources raise ``MetadataResourceNotFoundError``, and list reads + raise ``InvalidMetadataPageError`` unless ``offset >= 0`` and + ``1 <= limit <= 500``. + """ + + def close(self) -> None: + """Release the request's database session, even after a failed read.""" + ... def list_models( self, @@ -71,20 +82,26 @@ def list_models( *, offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataModel]: ... + ) -> MetadataPageResult[MetadataModel]: + """Return a page holding the selected catalog's single tax-benefit model.""" + ... def get_model( self, country_id: str, model_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataModel]: ... + ) -> MetadataDetailResult[MetadataModel]: + """Return the selected catalog's model when ``model_id`` is its UUID.""" + ... def get_model_by_country( self, country_id: str, policyengine_version: str | None = None, - ) -> MetadataModelSelectionResult: ... + ) -> MetadataModelSelectionResult: + """Return the country's model and the selected catalog's model version.""" + ... def list_model_versions( self, @@ -93,14 +110,18 @@ def list_model_versions( *, offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataModelVersionDetail]: ... + ) -> MetadataPageResult[MetadataModelVersionDetail]: + """Return a page holding the selected catalog's single model version.""" + ... def get_model_version( self, country_id: str, version_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataModelVersionDetail]: ... + ) -> MetadataDetailResult[MetadataModelVersionDetail]: + """Return the selected model version when ``version_id`` is its UUID.""" + ... def list_variables( self, @@ -110,14 +131,21 @@ def list_variables( offset: int = 0, limit: int = 100, search: str | None = None, - ) -> MetadataPageResult[MetadataVariable]: ... + ) -> MetadataPageResult[MetadataVariable]: + """Return variables in name order, optionally filtered by ``search``. + + ``search`` is a case-insensitive substring of the name, label or description. + """ + ... def get_variable( self, country_id: str, variable_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataVariable]: ... + ) -> MetadataDetailResult[MetadataVariable]: + """Return one variable of the selected catalog by UUID.""" + ... def list_parameters( self, @@ -127,14 +155,21 @@ def list_parameters( offset: int = 0, limit: int = 100, search: str | None = None, - ) -> MetadataPageResult[MetadataParameterSummary]: ... + ) -> MetadataPageResult[MetadataParameterSummary]: + """Return parameters in name order, optionally filtered by ``search``. + + ``search`` is a case-insensitive substring of the name, label or description. + """ + ... def get_parameter( self, country_id: str, parameter_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataParameterSummary]: ... + ) -> MetadataDetailResult[MetadataParameterSummary]: + """Return one parameter of the selected catalog by UUID.""" + ... def list_parameter_children( self, @@ -144,7 +179,12 @@ def list_parameter_children( parent_path: str = "", offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataParameterChild]: ... + ) -> MetadataPageResult[MetadataParameterChild]: + """Return the nodes and parameters one level below ``parent_path``. + + An empty ``parent_path`` lists the top level of the parameter tree. + """ + ... def list_parameter_values( self, @@ -156,14 +196,22 @@ def list_parameter_values( offset: int = 0, limit: int = 100, now: datetime | None = None, - ) -> MetadataPageResult[MetadataCanonicalParameterValue]: ... + ) -> MetadataPageResult[MetadataCanonicalParameterValue]: + """Return canonical values, optionally for one parameter. + + With ``current``, only values in effect on the UTC day of ``now`` (default: + the present) are returned. + """ + ... def get_parameter_value( self, country_id: str, value_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataCanonicalParameterValue]: ... + ) -> MetadataDetailResult[MetadataCanonicalParameterValue]: + """Return one canonical parameter value of the selected catalog by UUID.""" + ... def list_datasets( self, @@ -172,14 +220,18 @@ def list_datasets( *, offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataDataset]: ... + ) -> MetadataPageResult[MetadataDataset]: + """Return the logical input datasets of the selected catalog.""" + ... def get_dataset( self, country_id: str, dataset_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataDataset]: ... + ) -> MetadataDetailResult[MetadataDataset]: + """Return one logical input dataset of the selected catalog by UUID.""" + ... def list_regions( self, @@ -189,27 +241,41 @@ def list_regions( region_type: str | None = None, offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataRegion]: ... + ) -> MetadataPageResult[MetadataRegion]: + """Return the catalog's regions by code, optionally of one ``region_type``.""" + ... def get_region( self, country_id: str, region_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataRegion]: ... + ) -> MetadataDetailResult[MetadataRegion]: + """Return one region of the selected catalog by UUID.""" + ... def get_region_by_code( self, country_id: str, region_code: str, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataRegion]: ... + ) -> MetadataDetailResult[MetadataRegion]: + """Return one region by code; the national region's code is the country id.""" + ... def get_economy_options( self, country_id: str, policyengine_version: str | None = None, - ) -> MetadataEconomyOptionsResult: ... + ) -> MetadataEconomyOptionsResult: + """Return the options a client needs to configure an economy comparison. + + These are the current-law policy id, regions, years and datasets. + + Raise ``MetadataCatalogUnavailableError`` when the national region, its + dataset or the model version's options are missing from the catalog. + """ + ... class V2PolicyResourceService(Protocol): @@ -218,9 +284,20 @@ class V2PolicyResourceService(Protocol): def create_policy( self, command: NativePolicyCreationInput, - ) -> NativePolicyCreation: ... + ) -> NativePolicyCreation: + """Store the policy, or return the identical one already stored. + + Policies are immutable and deduplicated by canonical content hash; + ``created`` says whether this call inserted the policy. + """ + ... + + def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: + """Return one policy of the country with its parameter values. - def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: ... + Raise ``PolicyNotFoundError`` when no such policy exists in that country. + """ + ... def list_policies( self, @@ -229,7 +306,9 @@ def list_policies( tax_benefit_model_id: UUID | None = None, offset: int = 0, limit: int = 100, - ) -> PolicyPage: ... + ) -> PolicyPage: + """Return a page of the country's policies, optionally for one model.""" + ... class V2HouseholdResourceService(Protocol): @@ -238,14 +317,25 @@ class V2HouseholdResourceService(Protocol): def create_household( self, household_input: NativeHouseholdCreationInput, - ) -> NativeHouseholdCreation: ... + ) -> NativeHouseholdCreation: + """Store the household, or return the identical one already stored. + + Households are immutable and deduplicated by canonical content hash; + ``created`` says whether this call inserted the household. + """ + ... def get_household( self, *, country_id: str, household_id: UUID, - ) -> HouseholdRead: ... + ) -> HouseholdRead: + """Return one household of the country. + + Raise ``HouseholdNotFoundError`` when no such household exists there. + """ + ... def list_households( self, @@ -254,7 +344,9 @@ def list_households( default_year: int | None, offset: int, limit: int, - ) -> HouseholdPage: ... + ) -> HouseholdPage: + """Return a page of the country's households, optionally of one default year.""" + ... class V2UserPolicyResourceService(Protocol): @@ -263,14 +355,18 @@ class V2UserPolicyResourceService(Protocol): def create_user_policy( self, association_input: UserPolicyCreationInput, - ) -> UserPolicyRead: ... + ) -> UserPolicyRead: + """Save a policy for a user once both exist and their countries match.""" + ... def get_user_policy( self, *, country_id: str, association_id: UUID, - ) -> UserPolicyRead: ... + ) -> UserPolicyRead: + """Return the country's association, or raise ``UserPolicyNotFoundError``.""" + ... def list_user_policies( self, @@ -280,7 +376,9 @@ def list_user_policies( policy_id: UUID | None = None, offset: int = 0, limit: int = 100, - ) -> UserPolicyPage: ... + ) -> UserPolicyPage: + """Return a page of one user's associations, optionally for one policy.""" + ... def patch_user_policy( self, @@ -288,14 +386,18 @@ def patch_user_policy( country_id: str, association_id: UUID, association_input: UserPolicyUpdateInput, - ) -> UserPolicyRead: ... + ) -> UserPolicyRead: + """Update only the supplied presentation fields and return the association.""" + ... def delete_user_policy( self, *, country_id: str, association_id: UUID, - ) -> None: ... + ) -> None: + """Delete the country's association, or raise ``UserPolicyNotFoundError``.""" + ... class V2UserHouseholdResourceService(Protocol): @@ -304,14 +406,18 @@ class V2UserHouseholdResourceService(Protocol): def create_user_household( self, association_input: UserHouseholdCreationInput, - ) -> UserHouseholdRead: ... + ) -> UserHouseholdRead: + """Save a household for a user once both exist and their countries match.""" + ... def get_user_household( self, *, country_id: str, association_id: UUID, - ) -> UserHouseholdRead: ... + ) -> UserHouseholdRead: + """Return the country's association, or raise ``UserHouseholdNotFoundError``.""" + ... def list_user_households( self, @@ -321,7 +427,9 @@ def list_user_households( household_id: UUID | None = None, offset: int = 0, limit: int = 100, - ) -> UserHouseholdPage: ... + ) -> UserHouseholdPage: + """Return a page of one user's associations, optionally for one household.""" + ... def patch_user_household( self, @@ -329,20 +437,30 @@ def patch_user_household( country_id: str, association_id: UUID, association_input: UserHouseholdUpdateInput, - ) -> UserHouseholdRead: ... + ) -> UserHouseholdRead: + """Update only the supplied fields and return the association. + + A supplied ``household_id`` re-points the association to another + household, which must exist in the association's country. + """ + ... def delete_user_household( self, *, country_id: str, association_id: UUID, - ) -> None: ... + ) -> None: + """Delete the country's association, or raise ``UserHouseholdNotFoundError``.""" + ... class SimulationGatewayProbe(Protocol): """Minimal simulation-entrypoint health-check interface.""" - def health_check(self) -> bool: ... + def health_check(self) -> bool: + """Return whether the simulation entrypoint's health endpoint returns 200.""" + ... def _default_readiness_probe() -> bool: diff --git a/policyengine_api/fastapi_routes/health.py b/policyengine_api/fastapi_routes/health.py index b411512a8..6e0da7ec1 100644 --- a/policyengine_api/fastapi_routes/health.py +++ b/policyengine_api/fastapi_routes/health.py @@ -13,10 +13,14 @@ class HealthResponse(BaseModel): + """Health-check result showing that the API is up and serving requests.""" + status: Literal["healthy"] class SimulationGatewayHealthResponse(BaseModel): + """Health-check result showing that the API and its simulation gateway are up.""" + status: Literal["healthy"] simulation_gateway: Literal["healthy"] diff --git a/policyengine_api/fastapi_routes/metadata.py b/policyengine_api/fastapi_routes/metadata.py index 24a214a72..44bbe3f6d 100644 --- a/policyengine_api/fastapi_routes/metadata.py +++ b/policyengine_api/fastapi_routes/metadata.py @@ -26,6 +26,8 @@ class MetadataSuccessPayload(TypedDict): + """The v1 ``{status, message, result}`` envelope around a metadata document.""" + status: Literal["ok"] message: None result: JSONObject diff --git a/policyengine_api/fastapi_routes/responses.py b/policyengine_api/fastapi_routes/responses.py index 18ab9e249..114d51494 100644 --- a/policyengine_api/fastapi_routes/responses.py +++ b/policyengine_api/fastapi_routes/responses.py @@ -22,4 +22,5 @@ class LegacyJSONResponse(Response): media_type = "application/json" def render(self, content: JSONValue) -> bytes: + """Encode ``content`` byte-for-byte as the Flask metadata route does.""" return render_legacy_json(content) diff --git a/policyengine_api/fastapi_routes/v2/households/document_models.py b/policyengine_api/fastapi_routes/v2/households/document_models.py index d2b702b52..d64c01e9e 100644 --- a/policyengine_api/fastapi_routes/v2/households/document_models.py +++ b/policyengine_api/fastapi_routes/v2/households/document_models.py @@ -28,16 +28,31 @@ class StrictHouseholdDocumentModel(BaseModel): + """Base for household document records; rejects unknown keys and NaN or infinity.""" + model_config = ConfigDict(extra="forbid", allow_inf_nan=False) class HouseholdEntityRecord(StrictHouseholdDocumentModel): + """One group entity in a household, such as a tax unit or a benefit unit. + + ``id`` is unique within its collection and is what people's ``memberships`` + refer to. ``values`` maps variable names to values, and ``source_name`` + optionally keeps the entity's original name, such as its key in a v1 household. + """ + id: DocumentIdentifier source_name: DocumentIdentifier | None = None values: EntityValues class HouseholdPersonRecord(HouseholdEntityRecord): + """One person in a household. + + ``memberships`` maps each of the country's entity collections to the ``id`` + of the group the person belongs to, and must name every collection. + """ + memberships: dict[str, DocumentIdentifier] @@ -51,6 +66,12 @@ def _omit_spm_default(schema: dict[str, Any]) -> None: class USHouseholdDocument(StrictHouseholdDocumentModel): + """US household content: people and the US group entities they belong to. + + ``spm`` optionally selects an SPM forecast scenario and geography; omit it + rather than sending null. + """ + people: PersonRecords household: EntityRecords family: EntityRecords @@ -64,6 +85,7 @@ class USHouseholdDocument(StrictHouseholdDocumentModel): @field_validator("spm", mode="before") @classmethod def require_spm_object(cls, value: object) -> object: + """Reject an explicit null ``spm``; omit the field to make no selection.""" if value is None: raise ValueError("spm must be an object when supplied") return value @@ -72,6 +94,11 @@ def require_spm_object(cls, value: object) -> object: def serialize_document( self, handler: SerializerFunctionWrapHandler ) -> dict[str, Any]: + """Serialize the document, leaving out ``spm`` when no selection was made. + + Omitting the key keeps such documents identical to those stored before + ``spm`` existed. + """ document: dict[str, Any] = handler(self) if self.spm is None: document.pop("spm", None) @@ -79,6 +106,8 @@ def serialize_document( class UKHouseholdDocument(StrictHouseholdDocumentModel): + """UK household content: people, households and benefit units.""" + people: PersonRecords household: EntityRecords benunit: EntityRecords diff --git a/policyengine_api/fastapi_routes/v2/households/request_models.py b/policyengine_api/fastapi_routes/v2/households/request_models.py index 32619006a..80af28ab9 100644 --- a/policyengine_api/fastapi_routes/v2/households/request_models.py +++ b/policyengine_api/fastapi_routes/v2/households/request_models.py @@ -26,6 +26,7 @@ class HouseholdCreateRequest(BaseModel): @model_validator(mode="after") def require_country_document(self) -> HouseholdCreateRequest: + """Reject ``household_data`` whose collections belong to the other country.""" expected_type = ( USHouseholdDocument if self.country_id == "us" else UKHouseholdDocument ) @@ -34,6 +35,7 @@ def require_country_document(self) -> HouseholdCreateRequest: return self def to_service_input(self) -> NativeHouseholdCreationInput: + """Convert the request to the household service's creation input.""" return NativeHouseholdCreationInput.model_validate( self.model_dump(mode="python") ) diff --git a/policyengine_api/fastapi_routes/v2/households/response_models.py b/policyengine_api/fastapi_routes/v2/households/response_models.py index c51a8e5d6..35c645c9c 100644 --- a/policyengine_api/fastapi_routes/v2/households/response_models.py +++ b/policyengine_api/fastapi_routes/v2/households/response_models.py @@ -22,6 +22,8 @@ class StrictHouseholdAPIModel(BaseModel): + """Base for household responses; rejects unknown fields and NaN or infinity.""" + model_config = ConfigDict( extra="forbid", allow_inf_nan=False, @@ -30,6 +32,12 @@ class StrictHouseholdAPIModel(BaseModel): class HouseholdItem(StrictHouseholdAPIModel): + """A stored household. + + Households are immutable and deduplicated: creating one with identical + content returns the existing household and its ``id``. + """ + id: UUID country_id: CountryId default_year: DefaultYear | None @@ -39,14 +47,23 @@ class HouseholdItem(StrictHouseholdAPIModel): @classmethod def from_read(cls, item: HouseholdRead) -> HouseholdItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class HouseholdDetailResult(StrictHouseholdAPIModel): + """Result holding one household in ``item``.""" + item: HouseholdItem class HouseholdPageResult(StrictHouseholdAPIModel): + """One page of households, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[HouseholdItem] offset: int limit: int @@ -54,6 +71,7 @@ class HouseholdPageResult(StrictHouseholdAPIModel): @classmethod def from_page(cls, page: HouseholdPage) -> HouseholdPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[HouseholdItem.from_read(item) for item in page.items], offset=page.offset, @@ -66,17 +84,19 @@ def from_page(cls, page: HouseholdPage) -> HouseholdPageResult: class HouseholdSuccessResponse(StrictHouseholdAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT class HouseholdDetailResponse(HouseholdSuccessResponse[HouseholdDetailResult]): - pass + """Success response holding one household.""" class HouseholdPageResponse(HouseholdSuccessResponse[HouseholdPageResult]): - pass + """Success response holding one page of households.""" HOUSEHOLD_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { diff --git a/policyengine_api/fastapi_routes/v2/metadata/common.py b/policyengine_api/fastapi_routes/v2/metadata/common.py index bfe6fc90e..899243334 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/common.py +++ b/policyengine_api/fastapi_routes/v2/metadata/common.py @@ -68,6 +68,12 @@ def read_resource( response_type: type[ResponseT], operation: Callable[[V2MetadataResourceReader], object], ) -> ResponseT | JSONResponse: + """Run one metadata read and wrap its result in ``response_type``. + + Version, pagination and catalog failures become typed v2 error responses + (400, 404 or 503, and 500 for anything unexpected), and the reader is + closed whatever happens. + """ reader = None try: factory = dependencies.v2_metadata_reader_factory diff --git a/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py b/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py index cf3653a54..17ff62e27 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py @@ -31,6 +31,11 @@ def build_v2_metadata_geography_router( dependencies: NativeRouteDependencies, ) -> APIRouter: + """Build the v2 dataset, region and economy-option routes. + + Each request opens and closes its own metadata reader, so building the + router touches no database. + """ router = APIRouter(prefix="/v2", responses=ERROR_RESPONSES) @router.get( diff --git a/policyengine_api/fastapi_routes/v2/metadata/model_routes.py b/policyengine_api/fastapi_routes/v2/metadata/model_routes.py index 68336b1fe..927b52536 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/model_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/model_routes.py @@ -32,6 +32,11 @@ def build_v2_metadata_model_router( dependencies: NativeRouteDependencies, ) -> APIRouter: + """Build the v2 tax-benefit model, model-version and variable routes. + + Each request opens and closes its own metadata reader, so building the + router touches no database. + """ router = APIRouter(prefix="/v2", responses=ERROR_RESPONSES) @router.get( diff --git a/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py b/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py index 5683a3801..dfe11f152 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py @@ -30,6 +30,11 @@ def build_v2_metadata_parameter_router( dependencies: NativeRouteDependencies, ) -> APIRouter: + """Build the v2 parameter, parameter-tree and parameter-value routes. + + Each request opens and closes its own metadata reader, so building the + router touches no database. + """ router = APIRouter(prefix="/v2", responses=ERROR_RESPONSES) @router.get( diff --git a/policyengine_api/fastapi_routes/v2/metadata/response_models.py b/policyengine_api/fastapi_routes/v2/metadata/response_models.py index 6637998e9..d53660834 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/response_models.py +++ b/policyengine_api/fastapi_routes/v2/metadata/response_models.py @@ -24,6 +24,8 @@ class MetadataResourceSuccessResponse(StrictResponseModel, Generic[ResourceT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResourceT @@ -32,67 +34,73 @@ class MetadataResourceSuccessResponse(StrictResponseModel, Generic[ResourceT]): class MetadataModelPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataModel]] ): - pass + """Success response holding a page of the selected catalog's tax-benefit models. + + A catalog has one model per country, so only the first page has an item. + """ class MetadataModelDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataModel]] ): - pass + """Success response holding one tax-benefit model.""" class MetadataModelSelectionResponse( MetadataResourceSuccessResponse[MetadataModelSelectionResult] ): - pass + """Success response holding a country's model and its selected model version.""" class MetadataModelVersionPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataModelVersionDetail]] ): - pass + """Success response holding a page of the selected catalog's model versions. + + A catalog has one model version, so only the first page has an item. + """ class MetadataModelVersionDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataModelVersionDetail]] ): - pass + """Success response holding one model version.""" class MetadataVariablePageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataVariable]] ): - pass + """Success response holding one page of variables.""" class MetadataVariableDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataVariable]] ): - pass + """Success response holding one variable.""" class MetadataParameterPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataParameterSummary]] ): - pass + """Success response holding one page of parameters.""" class MetadataParameterDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataParameterSummary]] ): - pass + """Success response holding one parameter.""" class MetadataParameterChildPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataParameterChild]] ): - pass + """Success response holding one page of a parameter path's direct children.""" class MetadataParameterValuePageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataCanonicalParameterValue]] ): - pass + """Success response holding one page of parameter values.""" class MetadataParameterValueDetailResponse( @@ -100,34 +108,34 @@ class MetadataParameterValueDetailResponse( MetadataDetailResult[MetadataCanonicalParameterValue] ] ): - pass + """Success response holding one parameter value.""" class MetadataDatasetPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataDataset]] ): - pass + """Success response holding one page of datasets.""" class MetadataDatasetDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataDataset]] ): - pass + """Success response holding one dataset.""" class MetadataRegionPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataRegion]] ): - pass + """Success response holding one page of regions.""" class MetadataRegionDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataRegion]] ): - pass + """Success response holding one region.""" class MetadataEconomyOptionsResponse( MetadataResourceSuccessResponse[MetadataEconomyOptionsResult] ): - pass + """Success response holding the options for configuring an economy comparison.""" diff --git a/policyengine_api/fastapi_routes/v2/policies/response_models.py b/policyengine_api/fastapi_routes/v2/policies/response_models.py index 30cb2a242..0ae6b4887 100644 --- a/policyengine_api/fastapi_routes/v2/policies/response_models.py +++ b/policyengine_api/fastapi_routes/v2/policies/response_models.py @@ -25,6 +25,12 @@ class StrictPolicyAPIModel(BaseModel): class PolicyParameterValueItem(StrictPolicyAPIModel): + """One value a policy sets for a parameter from ``start_date``. + + ``parameter_name`` is the parameter's path in the model, and a null + ``end_date`` means the value has no end. + """ + id: UUID parameter_id: UUID parameter_name: str @@ -34,6 +40,12 @@ class PolicyParameterValueItem(StrictPolicyAPIModel): class PolicyItem(StrictPolicyAPIModel): + """A stored policy: parameter values for one tax-benefit model. + + Policies are immutable and deduplicated: creating one with identical + content returns the existing policy and its ``id``. + """ + id: UUID country_id: str tax_benefit_model_id: UUID @@ -43,14 +55,23 @@ class PolicyItem(StrictPolicyAPIModel): @classmethod def from_read(cls, item: PolicyRead) -> PolicyItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class PolicyDetailResult(StrictPolicyAPIModel): + """Result holding one policy in ``item``.""" + item: PolicyItem class PolicyPageResult(StrictPolicyAPIModel): + """One page of policies, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[PolicyItem] offset: int limit: int @@ -58,6 +79,7 @@ class PolicyPageResult(StrictPolicyAPIModel): @classmethod def from_page(cls, page: PolicyPage) -> PolicyPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[PolicyItem.from_read(item) for item in page.items], offset=page.offset, @@ -70,17 +92,19 @@ def from_page(cls, page: PolicyPage) -> PolicyPageResult: class PolicySuccessResponse(StrictPolicyAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT class PolicyDetailResponse(PolicySuccessResponse[PolicyDetailResult]): - pass + """Success response holding one policy.""" class PolicyPageResponse(PolicySuccessResponse[PolicyPageResult]): - pass + """Success response holding one page of policies.""" POLICY_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { diff --git a/policyengine_api/fastapi_routes/v2/user_households/response_models.py b/policyengine_api/fastapi_routes/v2/user_households/response_models.py index d269e541f..30ef20f76 100644 --- a/policyengine_api/fastapi_routes/v2/user_households/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_households/response_models.py @@ -19,10 +19,18 @@ class StrictUserHouseholdAPIModel(BaseModel): + """Base for user-household responses; rejects unknown fields.""" + model_config = ConfigDict(extra="forbid", from_attributes=True) class UserHouseholdItem(StrictUserHouseholdAPIModel): + """A household saved by a v2 user, with the user's name and description for it. + + Unlike households, associations can be renamed, pointed at another + household of the same country, or deleted. + """ + id: UUID country_id: CountryId user_id: UserId @@ -34,14 +42,23 @@ class UserHouseholdItem(StrictUserHouseholdAPIModel): @classmethod def from_read(cls, item: UserHouseholdRead) -> UserHouseholdItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class UserHouseholdDetailResult(StrictUserHouseholdAPIModel): + """Result holding one user-household association in ``item``.""" + item: UserHouseholdItem class UserHouseholdPageResult(StrictUserHouseholdAPIModel): + """One page of user-household associations, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[UserHouseholdItem] offset: int limit: int @@ -49,6 +66,7 @@ class UserHouseholdPageResult(StrictUserHouseholdAPIModel): @classmethod def from_page(cls, page: UserHouseholdPage) -> UserHouseholdPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[UserHouseholdItem.from_read(item) for item in page.items], offset=page.offset, @@ -61,6 +79,8 @@ def from_page(cls, page: UserHouseholdPage) -> UserHouseholdPageResult: class UserHouseholdSuccessResponse(StrictUserHouseholdAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT @@ -69,11 +89,11 @@ class UserHouseholdSuccessResponse(StrictUserHouseholdAPIModel, Generic[ResultT] class UserHouseholdDetailResponse( UserHouseholdSuccessResponse[UserHouseholdDetailResult] ): - pass + """Success response holding one user-household association.""" class UserHouseholdPageResponse(UserHouseholdSuccessResponse[UserHouseholdPageResult]): - pass + """Success response holding one page of user-household associations.""" USER_HOUSEHOLD_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { diff --git a/policyengine_api/fastapi_routes/v2/user_policies/response_models.py b/policyengine_api/fastapi_routes/v2/user_policies/response_models.py index 646e0bcf8..86c22abfa 100644 --- a/policyengine_api/fastapi_routes/v2/user_policies/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_policies/response_models.py @@ -25,6 +25,11 @@ class StrictUserPolicyAPIModel(BaseModel): class UserPolicyItem(StrictUserPolicyAPIModel): + """A policy saved by a v2 user, with the user's name and description for it. + + Unlike policies, associations can be renamed or deleted. + """ + id: UUID country_id: CountryId user_id: UserId @@ -36,14 +41,23 @@ class UserPolicyItem(StrictUserPolicyAPIModel): @classmethod def from_read(cls, item: UserPolicyRead) -> UserPolicyItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class UserPolicyDetailResult(StrictUserPolicyAPIModel): + """Result holding one user-policy association in ``item``.""" + item: UserPolicyItem class UserPolicyPageResult(StrictUserPolicyAPIModel): + """One page of user-policy associations, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[UserPolicyItem] offset: int limit: int @@ -51,6 +65,7 @@ class UserPolicyPageResult(StrictUserPolicyAPIModel): @classmethod def from_page(cls, page: UserPolicyPage) -> UserPolicyPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[UserPolicyItem.from_read(item) for item in page.items], offset=page.offset, @@ -63,17 +78,19 @@ def from_page(cls, page: UserPolicyPage) -> UserPolicyPageResult: class UserPolicySuccessResponse(StrictUserPolicyAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT class UserPolicyDetailResponse(UserPolicySuccessResponse[UserPolicyDetailResult]): - pass + """Success response holding one user-policy association.""" class UserPolicyPageResponse(UserPolicySuccessResponse[UserPolicyPageResult]): - pass + """Success response holding one page of user-policy associations.""" USER_POLICY_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { From 1536c5099a32b067bdd99371ad702487274eb4ae Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:23:52 +0400 Subject: [PATCH 08/12] Document the Flask routes, utilities and top-level modules Add module, function and class docstrings to the Flask blueprints, the payload validators and utilities, and the top-level application, constants, country, SPM, migration-flag and query-parameter modules. Reword the existing error-handler, user-profile, policy and cache-key docstrings into one-line imperative summaries, complete the Args sections of get_policy, set_policy and validate_country (whose entry described the wrapper's argument instead of the decorator's), and drop the update_user_profile promise of a 403 for auth0_id edits, which the code never makes. Only validator methods on the published SPM and query-parameter models gained docstrings; their class docstrings are unchanged. Co-Authored-By: Claude Opus 5.5 --- policyengine_api/api.py | 3 ++- policyengine_api/constants.py | 23 ++++++++++++++++--- policyengine_api/country.py | 16 +++++++++++++ policyengine_api/country_validation.py | 4 ++++ policyengine_api/migration_flags.py | 19 +++++++++++++++ policyengine_api/migration_registry.py | 6 +++++ policyengine_api/query_parameters.py | 2 ++ policyengine_api/routes/economy_routes.py | 15 ++++++++++++ policyengine_api/routes/error_routes.py | 16 +++++++------ policyengine_api/routes/home_routes.py | 2 ++ policyengine_api/routes/household_routes.py | 2 ++ policyengine_api/routes/metadata_routes.py | 2 ++ policyengine_api/routes/policy_routes.py | 21 ++++++++++------- .../routes/reform_impact_routes.py | 2 ++ .../routes/report_output_routes.py | 2 ++ policyengine_api/routes/simulation_routes.py | 2 ++ policyengine_api/routes/system_routes.py | 5 ++++ .../routes/user_profile_routes.py | 20 ++++++++++++---- policyengine_api/spm.py | 18 +++++++++++---- policyengine_api/utils/__init__.py | 2 ++ policyengine_api/utils/budget_window.py | 11 +++++++++ policyengine_api/utils/cache_utils.py | 7 +++--- policyengine_api/utils/deprecated_inputs.py | 1 + policyengine_api/utils/get_current_law.py | 7 ++++++ policyengine_api/utils/input_validation.py | 15 ++++++++++++ policyengine_api/utils/json.py | 14 +++++++++++ .../utils/payload_validators/__init__.py | 2 ++ .../payload_validators/validate_country.py | 13 ++++++++--- .../validate_household_payload.py | 2 ++ .../validate_set_policy_payload.py | 7 ++++++ policyengine_api/utils/singleton.py | 9 ++++++++ policyengine_api/utils/streaming_json.py | 2 ++ policyengine_api/worker_spm.py | 19 ++++++++++++++- policyengine_api/worker_spm_release.py | 6 +++++ 34 files changed, 263 insertions(+), 34 deletions(-) diff --git a/policyengine_api/api.py b/policyengine_api/api.py index ed2834da8..df4fea71c 100644 --- a/policyengine_api/api.py +++ b/policyengine_api/api.py @@ -1,5 +1,5 @@ # ruff: noqa: E402 -"""This is the main Flask app for the PolicyEngine API.""" +"""The PolicyEngine API Flask application, with its blueprints and extensions.""" import sys import time @@ -8,6 +8,7 @@ def log_timing(message: str) -> None: + """Print a startup milestone to stderr with the seconds since import began.""" elapsed = time.time() - start_time print(f"[{elapsed:.2f}s] {message}", file=sys.stderr, flush=True) diff --git a/policyengine_api/constants.py b/policyengine_api/constants.py index bcff670e5..b17a41b9f 100644 --- a/policyengine_api/constants.py +++ b/policyengine_api/constants.py @@ -1,3 +1,5 @@ +"""API version, country package versions and bundle metadata, resolved at import.""" + import hashlib import json from datetime import datetime @@ -49,6 +51,7 @@ def _resolve_distribution_version( def get_py_manifest() -> Any: # noqa: ANN401 # SimplePath, which open() rejects; a Path at runtime + """Return the path of the bundle manifest in the installed policyengine wheel.""" return distribution("policyengine").locate_file( "policyengine/data/bundle/manifest.json" ) @@ -89,6 +92,7 @@ def _resolve_bundle_package_versions(bundle: dict | None) -> dict[str, str]: def get_bundle_data_release(country_id: str) -> dict: + """Return the bundle's data release entry for a country, or {} if it has none.""" if not isinstance(_policyengine_bundle, dict): return {} @@ -104,6 +108,7 @@ def get_bundle_data_release(country_id: str) -> dict: def get_bundle_default_dataset(country_id: str) -> str | None: + """Return the bundle's default dataset name for a country, or None if unset.""" release = get_bundle_data_release(country_id) default_dataset = release.get("default_dataset") if default_dataset is None: @@ -116,6 +121,11 @@ def _bundle_dataset_display_label(default_dataset: object) -> str: def get_bundle_default_dataset_option(country_id: str) -> dict: + """Build the metadata dataset option for the bundle's certified default dataset. + + The option carries ``dataset_uri`` and ``data_version`` only when the bundle + records them. + """ release = get_bundle_data_release(country_id) default_dataset = release.get("default_dataset") label = _bundle_dataset_display_label(default_dataset) @@ -185,9 +195,10 @@ def get_bundle_default_dataset_option(country_id: str) -> dict: def _build_runtime_cache_version( scope: str, country_id: str, caller_version: str | None = None ) -> str: - """Build a compact version token for cache keys stored in legacy VARCHAR(10) - columns. The token changes whenever the relevant runtime or payload schema - changes, even if the country package version is unchanged. + """Build a compact cache version token that fits legacy VARCHAR(10) columns. + + The token changes whenever the relevant runtime or payload schema changes, + even if the country package version is unchanged. """ schema_version = str(RUNTIME_CACHE_SCHEMA_VERSIONS[scope]) prefix = "e" if scope == "economy_impact" else "r" @@ -215,10 +226,16 @@ def _build_runtime_cache_version( def get_economy_impact_cache_version( country_id: str, caller_version: str | None = None ) -> str: + """Return the cache version stored with economy impacts for this country. + + ``caller_version`` is the country model version the request asked for; it + defaults to the installed country package version. + """ return _build_runtime_cache_version("economy_impact", country_id, caller_version) def get_report_output_cache_version(country_id: str) -> str: + """Return the cache version that marks a report output current for this runtime.""" return _build_runtime_cache_version("report_output", country_id) diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 66c6ff415..ec0499113 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -1,3 +1,5 @@ +"""Loaded country tax-benefit systems, their metadata and household calculation.""" + import importlib import inspect import json @@ -121,7 +123,10 @@ def _record_calculation_failure( class PolicyEngineCountry: + """A country package's tax-benefit system with its prebuilt metadata document.""" + def __init__(self, country_package_name: str, country_id: str) -> None: + """Import the country package, build its tax-benefit system and its metadata.""" self.country_package_name = country_package_name self.country_id = country_id self.country_package = importlib.import_module(country_package_name) @@ -131,6 +136,7 @@ def __init__(self, country_package_name: str, country_id: str) -> None: self.build_metadata() def build_metadata(self) -> None: + """Build the JSON-safe metadata document served by the metadata endpoints.""" self.metadata = self._json_safe( dict( variables=self.build_variables(), @@ -172,6 +178,7 @@ def _json_safe( return value def build_microsimulation_options(self) -> dict: + """Build the regions, time periods and datasets offered for economy runs.""" # { region: [{ name: "uk", label: "the UK" }], time_period: [{ name: 2022, label: "2022", ... }] } options = dict() if self.country_id == "uk": @@ -334,6 +341,7 @@ def build_microsimulation_options(self) -> dict: return options def build_variables(self) -> dict: + """Describe every variable by name; enum variables list their values.""" variables = self.tax_benefit_system.variables variable_data = {} for variable_name, variable in variables.items(): @@ -369,6 +377,7 @@ def build_variables(self) -> dict: return variable_data def build_parameters(self) -> dict: + """Describe every parameter under ``gov``, with leaf parameters' values.""" parameters = self.tax_benefit_system.parameters parameter_data = {} for parameter in parameters.get_descendants(): @@ -428,6 +437,7 @@ def build_parameters(self) -> dict: return parameter_data def build_entities(self) -> dict: + """Describe every entity by key, with its roles (empty when it has none).""" data = {} for entity in self.tax_benefit_system.entities: entity_data = { @@ -695,6 +705,12 @@ def get_requested_computations( include_provided_values: bool = False, variable_names: set[str] | None = None, ) -> list[tuple[str, str, str, str]]: + """List the (entity plural, entity id, variable, period) cells to compute. + + A cell is requested when its value is None, or, with + ``include_provided_values``, whatever its value. ``variable_names`` keeps only + cells for those variables. + """ requested_computations = dpath.util.search( household, "*/*/*/*", diff --git a/policyengine_api/country_validation.py b/policyengine_api/country_validation.py index 28b08a280..a716ab60f 100644 --- a/policyengine_api/country_validation.py +++ b/policyengine_api/country_validation.py @@ -8,6 +8,8 @@ class CountryErrorPayload(TypedDict): + """The JSON error body returned for an unsupported country ID.""" + status: Literal["error"] message: str @@ -20,6 +22,7 @@ def __init__( country_id: str, available_country_ids: tuple[str, ...] = COUNTRIES, ) -> None: + """Record the rejected ID and the supported IDs listed in the message.""" self.country_id = country_id self.available_country_ids = available_country_ids super().__init__( @@ -28,6 +31,7 @@ def __init__( ) def to_payload(self) -> CountryErrorPayload: + """Return the ``{"status": "error", "message": ...}`` response body.""" return {"status": "error", "message": str(self)} diff --git a/policyengine_api/migration_flags.py b/policyengine_api/migration_flags.py index c1d5cbab3..67cf4f6bc 100644 --- a/policyengine_api/migration_flags.py +++ b/policyengine_api/migration_flags.py @@ -47,6 +47,8 @@ class RouteImplementation(StrEnum): @dataclass(frozen=True) class MigrationContext: + """The migration flag values in force for one request, as recorded in its logs.""" + route_group: str route_impl: RouteImplementation db_entity: str | None @@ -57,6 +59,7 @@ class MigrationContext: sim_compute: str | None def to_log_dict(self) -> dict: + """Return the flags as a flat log record, with ``route_impl`` as its value.""" return { "route_group": self.route_group, "route_impl": self.route_impl.value, @@ -109,6 +112,10 @@ def infer_route_group(path: str) -> str: def get_route_impl(route_group: str) -> RouteImplementation: + """Read ``ROUTE_IMPL_``, defaulting to the Flask fallback. + + Raises ValueError if the variable names an unknown implementation. + """ env_name = f"ROUTE_IMPL_{route_group.upper()}" return RouteImplementation( _read_choice( @@ -138,11 +145,19 @@ def from_environment(cls) -> RouteImplementationSettings: def get_db_write(entity: str) -> str: + """Read the ``DB_WRITE_`` write source, defaulting to Cloud SQL. + + Raises ValueError if the variable names an unknown source. + """ env_name = f"DB_WRITE_{entity.upper()}" return _read_choice(env_name, DEFAULT_DB_SOURCE, DB_WRITE_SOURCES) def get_db_read(entity: str) -> str: + """Read the ``DB_READ_`` read source, defaulting to Cloud SQL. + + Raises ValueError if the variable names an unknown source. + """ env_name = f"DB_READ_{entity.upper()}" return _read_choice(env_name, DEFAULT_DB_SOURCE, DB_READ_SOURCES) @@ -184,6 +199,10 @@ def get_v1_household_read_source() -> str: def get_sim_compute(flow: str) -> str: + """Read the ``SIM_COMPUTE_`` compute backend, defaulting to the old gateway. + + Raises ValueError if the variable names an unknown backend. + """ env_name = f"SIM_COMPUTE_{flow.upper()}" return _read_choice( env_name, diff --git a/policyengine_api/migration_registry.py b/policyengine_api/migration_registry.py index a99ae2131..ceb9a63b2 100644 --- a/policyengine_api/migration_registry.py +++ b/policyengine_api/migration_registry.py @@ -7,6 +7,12 @@ @dataclass(frozen=True) class RouteGroupConfig: + """A migratable route group and the URL path segments that identify its requests. + + ``db_entity`` and ``sim_flow`` select which database-source and + simulation-compute flags apply to the group's requests. + """ + name: str path_segments: tuple[str, ...] db_entity: str | None = None diff --git a/policyengine_api/query_parameters.py b/policyengine_api/query_parameters.py index 81c5c83e4..fc59d5d2c 100644 --- a/policyengine_api/query_parameters.py +++ b/policyengine_api/query_parameters.py @@ -173,6 +173,7 @@ class BudgetWindowEconomyQuery(EconomyQuery): @model_validator(mode="after") def validate_end_year(self) -> BudgetWindowEconomyQuery: + """Reject a window whose last year falls after BUDGET_WINDOW_MAX_END_YEAR.""" if int(self.start_year) + self.window_size - 1 > BUDGET_WINDOW_MAX_END_YEAR: raise ValueError( f"budget-window end_year must be {BUDGET_WINDOW_MAX_END_YEAR} or earlier" @@ -264,6 +265,7 @@ class DuplicateScalarQueryParameterError(ValueError): """Raised when a scalar query field is supplied more than once.""" def __init__(self, parameter_name: str) -> None: + """Record the public name of the repeated query parameter.""" self.parameter_name = parameter_name super().__init__( f"scalar query parameter {parameter_name!r} must not be repeated" diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py index cf493aac7..4aa952dec 100644 --- a/policyengine_api/routes/economy_routes.py +++ b/policyengine_api/routes/economy_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for society-wide economic impacts of a policy against a baseline.""" + import json from http import HTTPStatus @@ -73,6 +75,12 @@ def _bad_request_response(error: str | ValueError) -> Response: def get_economic_impact( country_id: str, policy_id: int, baseline_policy_id: int ) -> Response: + """Return the single-year economic impact, starting it if none is cached. + + A policy ID of 0 means current law. The result stays ``computing`` until the + simulation finishes. Invalid query parameters give 400 and a failed + calculation gives 502. + """ policy_id = int(policy_id or get_current_law_policy_id(country_id)) baseline_policy_id = int( baseline_policy_id or get_current_law_policy_id(country_id) @@ -122,6 +130,13 @@ def get_economic_impact( def get_budget_window_economic_impact( country_id: str, policy_id: int, baseline_policy_id: int ) -> Response: + """Return a multi-year budget-window impact with per-year progress. + + The body lists completed, computing and queued years, and the + ``X-PolicyEngine-Budget-Window-Cache`` header reports the cache outcome when + known. Invalid query parameters give 400; a failed calculation is reported + in the body with status 200. + """ policy_id = int(policy_id or get_current_law_policy_id(country_id)) baseline_policy_id = int( baseline_policy_id or get_current_law_policy_id(country_id) diff --git a/policyengine_api/routes/error_routes.py b/policyengine_api/routes/error_routes.py index b5f7ddb5f..ab6005351 100644 --- a/policyengine_api/routes/error_routes.py +++ b/policyengine_api/routes/error_routes.py @@ -1,3 +1,5 @@ +"""App-wide Flask error handlers that return the shared JSON error envelope.""" + from flask import Blueprint, Response from werkzeug.exceptions import ( HTTPException, @@ -10,41 +12,41 @@ @error_bp.app_errorhandler(404) def response_404(error: HTTPException) -> Response: - """Specific handler for 404 Not Found errors""" + """Return 404 Not Found as a JSON error body.""" return _make_error_response(error, 404, result=None) @error_bp.app_errorhandler(400) def response_400(error: HTTPException) -> Response: - """Specific handler for 400 Bad Request errors""" + """Return 400 Bad Request as a JSON error body.""" return _make_error_response(error, 400, result=None) @error_bp.app_errorhandler(401) def response_401(error: HTTPException) -> Response: - """Specific handler for 401 Unauthorized errors""" + """Return 401 Unauthorized as a JSON error body.""" return _make_error_response(error, 401, result=None) @error_bp.app_errorhandler(403) def response_403(error: HTTPException) -> Response: - """Specific handler for 403 Forbidden errors""" + """Return 403 Forbidden as a JSON error body.""" return _make_error_response(error, 403, result=None) @error_bp.app_errorhandler(500) def response_500(error: HTTPException) -> Response: - """Specific handler for 500 Internal Server errors""" + """Return 500 Internal Server Error as a JSON error body.""" return _make_error_response(error, 500, result=None) @error_bp.app_errorhandler(HTTPException) def response_http_exception(error: HTTPException) -> Response: - """Generic handler for HTTPException; should be raised if no specific handler is found""" + """Return any other HTTPException as a JSON error body with its own status.""" return _make_error_response(error, error.code, result=None) @error_bp.app_errorhandler(Exception) def response_generic_error(error: Exception) -> Response: - """Handler for any unhandled exceptions""" + """Return an unhandled exception as a 500 JSON error body carrying its message.""" return _make_error_response(error, 500, result=None) diff --git a/policyengine_api/routes/home_routes.py b/policyengine_api/routes/home_routes.py index ac05ccf44..b550d8ddb 100644 --- a/policyengine_api/routes/home_routes.py +++ b/policyengine_api/routes/home_routes.py @@ -1,3 +1,5 @@ +"""Flask route for the API's HTML landing page.""" + from flask import Blueprint home_bp = Blueprint("home", __name__) diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index 4d7ff2f13..0235fea20 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for v1 households and calculations, with optional copying to API v2.""" + import json import logging import time diff --git a/policyengine_api/routes/metadata_routes.py b/policyengine_api/routes/metadata_routes.py index f70e37bfd..ad658cbe8 100644 --- a/policyengine_api/routes/metadata_routes.py +++ b/policyengine_api/routes/metadata_routes.py @@ -1,3 +1,5 @@ +"""Flask route serving a country's model metadata document.""" + import json from flask import Blueprint, Response diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index 461da8c0e..5f5e4d61a 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for v1 policies and saved user policies, optionally mirrored to v2.""" + import json import time @@ -142,15 +144,15 @@ def _serialize_policy(policy: Policy) -> dict: @policy_bp.route("//policy/", methods=["GET"]) @validate_country def get_policy(country_id: str, policy_id: int | str) -> Response: - """Get policy data for a given country and policy ID. + """Return one stored v1 policy. Args: - country_id (str) - policy_id (int | str) + country_id (str): Country whose policies are searched. + policy_id (int | str): ID of the stored policy; cast to int. Returns: - Response: A Flask response object containing the - policy data in JSON format + Response: The policy as JSON under ``result``. An unknown ID gives 404, + and an invalid policy read-source setting gives 503. """ try: @@ -176,11 +178,14 @@ def get_policy(country_id: str, policy_id: int | str) -> Response: @policy_bp.route("//policy", methods=["POST"]) @validate_country def set_policy(country_id: str) -> Response: - """Set policy data for given country and policy. If policy already exists, - return existing policy and 200. + """Store a policy, or return the ID of an identical one already stored. + + A policy is identical when its country, content hash and label all match; + that case returns 200, a new policy returns 201. An invalid payload gives + 400, and an unavailable write configuration or v2 mirror gives 503. Args: - country_id (str) + country_id (str): Country to store the policy under. """ payload = request.json diff --git a/policyengine_api/routes/reform_impact_routes.py b/policyengine_api/routes/reform_impact_routes.py index 82b2a9f46..21bbf6df9 100644 --- a/policyengine_api/routes/reform_impact_routes.py +++ b/policyengine_api/routes/reform_impact_routes.py @@ -1,3 +1,5 @@ +"""Flask route listing recent cached reform impact calculations.""" + import json from datetime import datetime from typing import Any diff --git a/policyengine_api/routes/report_output_routes.py b/policyengine_api/routes/report_output_routes.py index ea188dcc4..b8a919de3 100644 --- a/policyengine_api/routes/report_output_routes.py +++ b/policyengine_api/routes/report_output_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for creating, reading and updating v1 report outputs.""" + import json import jsonschema diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index f0853c53b..76ca6e7db 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for creating, reading and updating v1 simulation records.""" + import json import jsonschema diff --git a/policyengine_api/routes/system_routes.py b/policyengine_api/routes/system_routes.py index 450459083..67b7959d1 100644 --- a/policyengine_api/routes/system_routes.py +++ b/policyengine_api/routes/system_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for the liveness and readiness probes and the legacy OpenAPI spec.""" + from flask import Blueprint, Response, jsonify from policyengine_api.readiness import is_ready @@ -8,11 +10,13 @@ @system_bp.route("/liveness-check", methods=["GET"]) def liveness_check() -> Response: + """Return 200 "OK" whenever the process can serve a request, even during warmup.""" return Response("OK", status=200, headers={"Content-Type": "text/plain"}) @system_bp.route("/readiness-check", methods=["GET"]) def readiness_check() -> Response: + """Return 200 "OK" once startup warmup has finished, and 503 "NOT READY" before.""" # The service is not ready until startup warmup compiles the simulation # machinery. Liveness remains unconditional so the worker is not restarted. if not is_ready(): @@ -26,4 +30,5 @@ def readiness_check() -> Response: @system_bp.route("/specification", methods=["GET"]) def get_specification() -> Response: + """Return the legacy public OpenAPI specification as JSON.""" return jsonify(OPENAPI_SPECIFICATION) diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index 8bcd68bd7..4ffcccbad 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for creating, reading and updating v1 user profiles.""" + import json from flask import Blueprint, Response, request @@ -30,7 +32,11 @@ def _serialize_user_profile( @user_profile_bp.route("//user-profile", methods=["POST"]) @validate_country def set_user_profile(country_id: str) -> Response: - """Creates a new user_profile""" + """Create a user profile for an Auth0 ID, or return the one it already has. + + Returns 201 for a new profile and 200 for an existing one; the Auth0 ID is + never echoed back. + """ payload = request.json if payload is None: raise BadRequest("Payload missing from request") @@ -62,6 +68,11 @@ def set_user_profile(country_id: str) -> Response: @user_profile_bp.route("//user-profile", methods=["GET"]) @validate_country def get_user_profile(country_id: str) -> Response: + """Return a user profile looked up by ``auth0_id`` or, failing that, ``user_id``. + + The Auth0 ID is included only when the caller looked the profile up by it. + Neither identifier gives 400 and an unknown user gives 404. + """ auth0_id = request.args.get("auth0_id") user_id = request.args.get("user_id") @@ -101,9 +112,10 @@ def get_user_profile(country_id: str) -> Response: @user_profile_bp.route("//user-profile", methods=["PUT"]) @validate_country def update_user_profile(country_id: str) -> Response: - """Update any part of a user_profile, given a user_id, - except the auth0_id value; any attempt to edit this - will assume malicious intent and 403 + """Update the username, primary country or start date of profile ``user_id``. + + The Auth0 ID cannot be changed: an ``auth0_id`` in the payload is ignored. + A null ``user_id`` gives 400 and an unknown one gives 404. """ payload = request.json diff --git a/policyengine_api/spm.py b/policyengine_api/spm.py index ddfd78002..0222b472a 100644 --- a/policyengine_api/spm.py +++ b/policyengine_api/spm.py @@ -93,12 +93,14 @@ def serialize_selection( # noqa: ANN201 # pydantic would build the output sche @field_validator("as_of") @classmethod def validate_as_of(cls, value: str | None) -> str | None: + """Reject an ``as_of`` that is not an exact ISO date such as 2025-01-31.""" if value is not None and date.fromisoformat(value).isoformat() != value: raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)") return value @model_validator(mode="after") def validate_location(self) -> Self: + """Require a non-blank ``geography_id`` for a metro; forbid it otherwise.""" if self.geography_kind == "metro": if not self.geography_id or not self.geography_id.strip(): raise ValueError("An SPM area selection requires geography_id") @@ -121,12 +123,14 @@ class SPMResolvedConfiguration(BaseModel): @field_validator("as_of") @classmethod def validate_as_of(cls, value: str | None) -> str | None: + """Reject an ``as_of`` that is not an exact ISO date such as 2025-01-31.""" if value is not None and date.fromisoformat(value).isoformat() != value: raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)") return value @model_validator(mode="after") def validate_location(self) -> SPMResolvedConfiguration: + """Require a non-blank ``geography_id`` for a metro; forbid it otherwise.""" if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): raise ValueError("An SPM area selection requires geography_id") @@ -153,6 +157,7 @@ class SPMProvenance(BaseModel): @field_validator("years") @classmethod def validate_years(cls, years: tuple[str, ...]) -> tuple[str, ...]: + """Require sorted, unique four-digit year strings.""" if any( len(year) != 4 or not year.isascii() or not year.isdigit() for year in years ): @@ -163,6 +168,7 @@ def validate_years(cls, years: tuple[str, ...]) -> tuple[str, ...]: @model_validator(mode="after") def validate_location(self) -> SPMProvenance: + """Require a non-blank ``geography_id`` on a metro receipt, else forbid it.""" if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): raise ValueError("A metro SPM receipt requires geography_id") @@ -202,6 +208,7 @@ class SPMComparisonProvenance(BaseModel): @model_validator(mode="after") def validate_compatible_receipts(self) -> SPMComparisonProvenance: + """Reject a comparison whose baseline and reform used different SPM receipts.""" if self.baseline.receipt != self.reform.receipt: raise ValueError("baseline and reform SPM receipts differ") return self @@ -321,11 +328,13 @@ class SPMValidationError(ValueError): """An SPM request/configuration error safe for the API's 4xx envelope.""" def __init__(self, code: str, message: str) -> None: + """Record a stable machine-readable error code and the client-facing message.""" self.code = code self.message = message super().__init__(message) def to_dict(self) -> dict[str, str]: + """Return the ``{"code", "message"}`` entry for a response's ``errors`` list.""" return {"code": self.code, "message": self.message} @@ -580,11 +589,12 @@ def normalize_spm_selection( def _validate_metro_selection(forecast: SPMForecast, selection: SPMSelection) -> None: - """A selected area must exist in the artifact, without guessing a year. + """Raise SPM_GEOGRAPHY_UNAVAILABLE unless the selected metro is in the artifact. - Area definitions may change between years. Accept an area available in any - covered year; the country owns validation for the actual calculation year. - This reads geography metadata and creates no household calculation receipt. + Area definitions may change between years, so this does not guess a year: it + accepts an area available in any covered year, and the country owns + validation for the actual calculation year. This reads geography metadata + and creates no household calculation receipt. """ for year in forecast.years: try: diff --git a/policyengine_api/utils/__init__.py b/policyengine_api/utils/__init__.py index 1aeb79c92..85ca5a173 100644 --- a/policyengine_api/utils/__init__.py +++ b/policyengine_api/utils/__init__.py @@ -1,3 +1,5 @@ +"""Shared helpers for caching, JSON values, current-law lookup and singletons.""" + from .cache_utils import * # noqa: F403 from .get_current_law import get_current_law_policy_id as get_current_law_policy_id from .json import * # noqa: F403 diff --git a/policyengine_api/utils/budget_window.py b/policyengine_api/utils/budget_window.py index 6acbd5b3a..d827c66eb 100644 --- a/policyengine_api/utils/budget_window.py +++ b/policyengine_api/utils/budget_window.py @@ -1,3 +1,5 @@ +"""Limits and request normalisation for multi-year budget-window calculations.""" + from dataclasses import dataclass from typing import Literal @@ -8,6 +10,8 @@ @dataclass(frozen=True) class BudgetWindowRequestSetup: + """A validated budget-window request, expanded into its years and cache period.""" + start_year: str window_size: int years: list[str] @@ -15,11 +19,13 @@ class BudgetWindowRequestSetup: def build_budget_window_years(*, start_year: str, window_size: int) -> list[str]: + """Return ``window_size`` consecutive years from ``start_year``, as strings.""" start_year_int = int(start_year) return [str(start_year_int + index) for index in range(window_size)] def build_budget_window_time_period(*, start_year: str, window_size: int) -> str: + """Return the ``budget_window::`` period naming the whole window.""" return f"budget_window:{start_year}:{window_size}" @@ -29,6 +35,11 @@ def build_budget_window_request_setup( window_size: int, target: Literal["general", "cliff"], ) -> BudgetWindowRequestSetup: + """Validate a budget-window request and expand it into its years. + + Raises ValueError for a target other than ``general``, a window outside + 1..BUDGET_WINDOW_MAX_YEARS, or a last year after BUDGET_WINDOW_MAX_END_YEAR. + """ if target != "general": raise ValueError("Budget-window calculations only support target='general'") diff --git a/policyengine_api/utils/cache_utils.py b/policyengine_api/utils/cache_utils.py index e977962f4..c5f8bf516 100644 --- a/policyengine_api/utils/cache_utils.py +++ b/policyengine_api/utils/cache_utils.py @@ -10,9 +10,10 @@ def make_cache_key(*args: object, **kwargs: object) -> str: # pylint: disable=unused-argument - """Make a hash to uniquely identify a cache entry. - keep it fast, adding overhead to try to add some minor chance of a - cache hit is not worth it. + """Hash the current request's full path and JSON or form body into a cache key. + + The arguments are ignored. Keep it fast: adding overhead for a minor extra + chance of a cache hit is not worth it. Use a cryptographic digest (SHA-256) rather than the builtin `hash()`, whose output depends on PYTHONHASHSEED and is therefore diff --git a/policyengine_api/utils/deprecated_inputs.py b/policyengine_api/utils/deprecated_inputs.py index 1e1af7314..85cf7ee0b 100644 --- a/policyengine_api/utils/deprecated_inputs.py +++ b/policyengine_api/utils/deprecated_inputs.py @@ -38,6 +38,7 @@ class DeprecatedVariableWarning: @property def message(self) -> str: + """The user-facing warning naming the ignored input, its location and hint.""" location = f"`{self.entity_plural}/{self.entity_id}`" if self.entity_plural == "axes": location = f"`axes[{self.entity_id}].name`" diff --git a/policyengine_api/utils/get_current_law.py b/policyengine_api/utils/get_current_law.py index 929990461..1853570fc 100644 --- a/policyengine_api/utils/get_current_law.py +++ b/policyengine_api/utils/get_current_law.py @@ -1,4 +1,11 @@ +"""Lookup of each country's built-in current-law policy ID.""" + + def get_current_law_policy_id(country_id: str) -> int: + """Return the ID of the stored policy that represents current law in a country. + + Raises KeyError for a country without one (Israel has none here). + """ return { "uk": 1, "us": 2, diff --git a/policyengine_api/utils/input_validation.py b/policyengine_api/utils/input_validation.py index 3029c2760..af691f508 100644 --- a/policyengine_api/utils/input_validation.py +++ b/policyengine_api/utils/input_validation.py @@ -5,6 +5,11 @@ @dataclass(frozen=True) class UnrecognizedInput: + """One payload name the country model does not recognise, and where it appears. + + The entity plurals are set only for a variable placed on the wrong entity. + """ + input_type: str name: str path: str @@ -13,6 +18,7 @@ class UnrecognizedInput: @property def message(self) -> str: + """A one-sentence, user-facing description of the problem and where it is.""" if self.input_type == "household_entity": return ( f"Unrecognized household entity group `{self.name}` at `{self.path}`." @@ -32,6 +38,7 @@ def message(self) -> str: return f"Unrecognized household variable `{self.name}` at `{self.path}`." def to_dict(self) -> dict: + """Return the JSON error entry, omitting entity plurals that are not set.""" data = { "type": self.input_type, "name": self.name, @@ -61,6 +68,12 @@ def find_unrecognized_household_inputs( household: dict, metadata: dict, ) -> list[UnrecognizedInput]: + """Find unknown entity groups, variables and axis variables in a household. + + Also reports a known variable placed on the wrong entity. Role membership + keys (such as ``members``) are not variables and are skipped, as are + sections that are not objects. + """ if not isinstance(household, dict): return [] @@ -134,6 +147,7 @@ def find_unrecognized_policy_inputs( policy: dict, metadata: dict, ) -> list[UnrecognizedInput]: + """Find policy keys that are not parameters in the country metadata.""" if not isinstance(policy, dict): return [] @@ -150,6 +164,7 @@ def find_unrecognized_policy_inputs( def format_unrecognized_inputs_message(errors: list[UnrecognizedInput]) -> str: + """Join the problems' messages into one error message for the response.""" return "Unrecognized calculate input(s): " + "; ".join( error.message for error in errors ) diff --git a/policyengine_api/utils/json.py b/policyengine_api/utils/json.py index 54f54e8d7..668803ef8 100644 --- a/policyengine_api/utils/json.py +++ b/policyengine_api/utils/json.py @@ -1,3 +1,5 @@ +"""Content hashing and JSON-safe conversion of policy, household and metadata values.""" + from __future__ import annotations import base64 @@ -13,6 +15,10 @@ def make_hashable( o: object, ) -> Any: # noqa: ANN401 # returns scalars unchanged, and callers sort them + """Convert nested lists, dicts and sets into tuples, sorting dict items and sets. + + Equal values therefore convert identically whatever their key or member order. + """ if isinstance(o, (tuple, list)): return tuple(make_hashable(e) for e in o) @@ -26,12 +32,20 @@ def make_hashable( def hash_object(o: object) -> str: + """Return a base64 SHA-256 digest of a value that ignores dict key order. + + It is the stored policy and household hash, and keys cached calculations. + """ hasher = hashlib.sha256() hasher.update(repr(make_hashable(o)).encode()) return base64.b64encode(hasher.digest()).decode() def get_safe_json(value: object) -> JSONValue: + """Return a JSON-safe copy of a metadata value, writing infinities as strings. + + Values of any other type (dates, None, objects) become None. + """ # This function is used when constructing metadata, # which will be consumed by JS, hence Infinity, not .inf if isinstance(value, (int, float)): diff --git a/policyengine_api/utils/payload_validators/__init__.py b/policyengine_api/utils/payload_validators/__init__.py index 6da6b53cb..ba014323d 100644 --- a/policyengine_api/utils/payload_validators/__init__.py +++ b/policyengine_api/utils/payload_validators/__init__.py @@ -1,3 +1,5 @@ +"""Validators for Flask route inputs: country IDs and policy and household payloads.""" + from .validate_country import validate_country as validate_country from .validate_household_payload import ( validate_household_payload as validate_household_payload, diff --git a/policyengine_api/utils/payload_validators/validate_country.py b/policyengine_api/utils/payload_validators/validate_country.py index 1ab713804..33f2e7a15 100644 --- a/policyengine_api/utils/payload_validators/validate_country.py +++ b/policyengine_api/utils/payload_validators/validate_country.py @@ -1,3 +1,5 @@ +"""Flask view decorator that rejects unsupported country IDs.""" + from __future__ import annotations from functools import wraps @@ -20,13 +22,18 @@ def validate_country( func: Callable[..., _ViewResult], ) -> Callable[..., _ViewResult | Response]: - """Validate that a country ID is valid. If not, return a 400 response. + """Wrap a Flask view so an unsupported ``country_id`` gets 400 before the view runs. + + The 400 body is the shared JSON error payload, sent with the legacy v1 + content type that native FastAPI routes must match. Args: - country_id (str): The country ID to validate. + func (Callable): The view to wrap; it takes ``country_id`` as its first + argument. Returns: - Response(400) if country is not valid, else continues + Callable: The wrapped view, which returns a 400 response for an + unsupported country and the view's own result otherwise. """ diff --git a/policyengine_api/utils/payload_validators/validate_household_payload.py b/policyengine_api/utils/payload_validators/validate_household_payload.py index 73da22c84..750c9874a 100644 --- a/policyengine_api/utils/payload_validators/validate_household_payload.py +++ b/policyengine_api/utils/payload_validators/validate_household_payload.py @@ -1,3 +1,5 @@ +"""Validation of the body of a v1 household create request.""" + from __future__ import annotations from typing import TYPE_CHECKING diff --git a/policyengine_api/utils/payload_validators/validate_set_policy_payload.py b/policyengine_api/utils/payload_validators/validate_set_policy_payload.py index 6bb9002fd..9c9d2f098 100644 --- a/policyengine_api/utils/payload_validators/validate_set_policy_payload.py +++ b/policyengine_api/utils/payload_validators/validate_set_policy_payload.py @@ -1,4 +1,11 @@ +"""Validation of the body of a v1 policy create request.""" + + def validate_set_policy_payload(payload: dict) -> tuple[bool, str | None]: + """Check that a policy body has a ``data`` object and an optional string ``label``. + + Return ``(True, None)`` when valid, or ``(False, message)`` naming the problem. + """ # Check that all required keys are present required_keys = ["data"] missing_keys = [key for key in required_keys if key not in payload] diff --git a/policyengine_api/utils/singleton.py b/policyengine_api/utils/singleton.py index ecba9af47..bc801566f 100644 --- a/policyengine_api/utils/singleton.py +++ b/policyengine_api/utils/singleton.py @@ -1,7 +1,16 @@ +"""Metaclass for classes that have a single shared instance per process.""" + + class Singleton(type): + """Metaclass whose classes build one instance and return it on every later call.""" + _instances = {} def __call__(cls, *args: object, **kwargs: object) -> object: + """Return the shared instance, creating it on the first call. + + Arguments to later calls are ignored. + """ if cls not in cls._instances: cls._instances[cls] = super().__call__(*args, **kwargs) return cls._instances[cls] diff --git a/policyengine_api/utils/streaming_json.py b/policyengine_api/utils/streaming_json.py index f0589a525..488771a09 100644 --- a/policyengine_api/utils/streaming_json.py +++ b/policyengine_api/utils/streaming_json.py @@ -20,11 +20,13 @@ def should_stream_body(body: bytes) -> bool: + """Return True when a body is large enough that it must be sent chunked.""" return len(body) >= STREAMING_THRESHOLD_BYTES def iter_body_chunks( body: bytes, chunk_size: int = STREAMING_CHUNK_BYTES ) -> Iterator[bytes]: + """Yield a body in fixed-size slices for a chunked response.""" for start in range(0, len(body), chunk_size): yield body[start : start + chunk_size] diff --git a/policyengine_api/worker_spm.py b/policyengine_api/worker_spm.py index 498b5c5af..f0e446830 100644 --- a/policyengine_api/worker_spm.py +++ b/policyengine_api/worker_spm.py @@ -28,7 +28,9 @@ def get_spm_capability( /, *, policyengine_version: str | None, - ) -> object: ... + ) -> object: + """Return the worker's SPM capability document, or None if it has none.""" + ... def validate_worker_spm( @@ -39,6 +41,13 @@ def validate_worker_spm( policyengine_version: str | None = None, model_version: str | None = None, ) -> dict | None: + """Resolve an SPM selection and confirm the selected worker can execute it. + + Return the resolved selection, or None for a country without SPM support. + Raise SPMValidationError (SPM_CONFIGURATION_UNAVAILABLE) + when the worker advertises no canonical SPM capability, a different + contract, or a different forecast artifact from the API's. + """ country_id = country_id.lower() if isinstance(country_id, str) else country_id resolved = normalize_spm_selection(country_id, selection) if resolved is None: @@ -79,6 +88,14 @@ def validate_worker_result( expected_year: str | int | None = None, expected_years: list[str] | None = None, ) -> None: + """Check that a worker's economy result carries SPM receipts matching the request. + + An annual result must hold a comparison receipt whose baseline and reform + resolve to ``selection`` and cover ``expected_year``; a budget-window result + must have one valid annual result per requested year. Raise + SPMValidationError (SPM_CONFIGURATION_UNAVAILABLE) on any mismatch. Do nothing + when no selection was requested. + """ if selection is None: return try: diff --git a/policyengine_api/worker_spm_release.py b/policyengine_api/worker_spm_release.py index 17d430325..d5bee00c8 100644 --- a/policyengine_api/worker_spm_release.py +++ b/policyengine_api/worker_spm_release.py @@ -22,6 +22,7 @@ class RegistryCapability: """Expose one registry document through the request validator interface.""" def __init__(self, registry: dict) -> None: + """Wrap a parsed worker registry document.""" self.registry = registry def get_spm_capability( @@ -31,6 +32,11 @@ def get_spm_capability( *, policyengine_version: str | None, ) -> object: + """Return the registry's SPM capability for ``policyengine_version``, or None. + + The registry keys capabilities by PolicyEngine.py version alone, so + ``country`` and ``model_version`` are not consulted. + """ capabilities = self.registry.get("spm_capabilities") if not isinstance(capabilities, dict): return None From ab1435bf4be7a0b8becbb042f4450f2683308e44 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:23:59 +0400 Subject: [PATCH 09/12] Document the v1 services Add docstrings to the v1 economy, household, policy, report, simulation and user services, their legacy-to-v2 mirroring helpers and the budget- window cache, including their Pydantic option, result and spec models. They say what a reader cannot see from the signature: which results are cached and how, which writes are mirrored to v2, what None means, and which errors are raised. Reword the existing summaries that ran on without a blank line or past 88 columns, correct get_economic_impact, which claimed to return a tuple, and replace the pass bodies of two exception classes with their new docstrings. Co-Authored-By: Claude Opus 5.5 --- policyengine_api/services/__init__.py | 2 + .../services/budget_window_cache.py | 34 +++++++++++ .../services/economy_comparison.py | 57 +++++++++++++++++- policyengine_api/services/economy_service.py | 59 ++++++++++++++----- .../services/household_calculation_service.py | 31 +++++++++- .../services/household_mirroring.py | 4 +- .../services/household_service.py | 21 +++++++ policyengine_api/services/metadata_service.py | 5 ++ policyengine_api/services/policy_mirroring.py | 4 +- policyengine_api/services/policy_service.py | 19 ++++++ .../services/reform_impacts_service.py | 26 ++++++++ .../services/report_output_alias_service.py | 15 +++++ .../services/report_output_service.py | 33 +++++++++++ .../services/report_run_service.py | 17 ++++++ .../services/report_spec_service.py | 32 ++++++++++ .../services/simulation_run_service.py | 17 ++++++ .../services/simulation_service.py | 21 +++++++ .../services/simulation_spec_service.py | 17 ++++++ .../services/user_policy_mirroring.py | 13 +++- .../services/user_policy_service.py | 27 +++++++++ policyengine_api/services/user_service.py | 15 +++++ 21 files changed, 445 insertions(+), 24 deletions(-) diff --git a/policyengine_api/services/__init__.py b/policyengine_api/services/__init__.py index ca7322caf..5bd08f8f8 100644 --- a/policyengine_api/services/__init__.py +++ b/policyengine_api/services/__init__.py @@ -1 +1,3 @@ +"""Business logic for the v1 API resources; ``v2`` holds the v2 resource services.""" + # from .economy_service import EconomyService, JobService, PolicyService diff --git a/policyengine_api/services/budget_window_cache.py b/policyengine_api/services/budget_window_cache.py index 1ec50c10d..edff6530b 100644 --- a/policyengine_api/services/budget_window_cache.py +++ b/policyengine_api/services/budget_window_cache.py @@ -50,10 +50,12 @@ class BudgetWindowCacheState(BaseModel): error: dict[str, Any] | None = None def to_payload(self) -> dict[str, Any]: + """Return the state as the cache JSON payload, omitting unset fields.""" return self.model_dump(exclude_none=True) @model_validator(mode="after") def require_fields_for_status(self) -> BudgetWindowCacheState: + """Raise ValueError unless every field this status depends on is set.""" missing = [ field for field in self._required_fields[self.status] @@ -65,6 +67,7 @@ def require_fields_for_status(self) -> BudgetWindowCacheState: @classmethod def from_payload(cls, payload: object) -> BudgetWindowCacheState | None: + """Parse a decoded cache payload, or return None if it is not a valid state.""" try: return cls.model_validate(payload) except ValidationError: @@ -79,6 +82,7 @@ def __init__( client: CacheBackend | None = None, namespace: CacheNamespace | None = None, ) -> None: + """Bind to a cache client and namespace, defaulting to the runtime cache.""" if client is None or namespace is None: context = get_runtime_cache_context() client = client or context.client @@ -102,6 +106,11 @@ def build_key( options_hash: str | None, api_version: str, ) -> str: + """Return the cache key for one budget-window report request. + + Every argument is part of the key, so a change to any of them (including the + API version or the options hash) starts a separate report. + """ return self.namespace.key( BUDGET_WINDOW_CACHE_FAMILY, BUDGET_WINDOW_CACHE_SCHEMA_VERSION, @@ -202,6 +211,12 @@ def claim_batch_start( claim_token: str, observability_id: str | None, ) -> bool: + """Atomically claim the right to submit this report's batch job. + + Return True only if no state existed for the key, in which case a short-lived + "starting" state now holds the claim. Raise CacheCoordinationError if the cache + is unreachable, so callers never submit duplicate work blindly. + """ state = BudgetWindowCacheState( status="starting", observability_id=observability_id, @@ -219,6 +234,11 @@ def clear_starting_claim( claim_token: str, observability_id: str | None, ) -> None: + """Release this caller's "starting" claim so another request can submit. + + The claim is deleted only if it still holds this token; cache errors are + swallowed because the claim expires on its own. + """ state = BudgetWindowCacheState( status="starting", observability_id=observability_id, @@ -238,6 +258,11 @@ def store_submitted( batch_job_id: str, observability_id: str | None, ) -> None: + """Replace the starting claim with the submitted batch job ID. + + Raise CacheCoordinationError if the state cannot be written, since later polls + would otherwise lose track of the running batch. + """ state = BudgetWindowCacheState( status="submitted", observability_id=observability_id, @@ -281,6 +306,7 @@ def set_completed_result( result: dict[str, Any], observability_id: str | None, ) -> bool: + """Store the completed report result; return False if the cache write failed.""" return self._set_recoverable_state( cache_key, BudgetWindowCacheState( @@ -297,6 +323,10 @@ def set_terminal_error( error: dict[str, str], observability_id: str | None, ) -> bool: + """Store an SPM validation failure so later requests replay it, not resubmit. + + Return False if the cache write failed. + """ return self._set_recoverable_state( cache_key, BudgetWindowCacheState( @@ -314,6 +344,10 @@ def set_execution_failure( result: dict[str, Any], observability_id: str | None, ) -> bool: + """Store a failed batch result so later requests return it without resubmitting. + + Return False if the cache write failed. + """ return self._set_recoverable_state( cache_key, BudgetWindowCacheState( diff --git a/policyengine_api/services/economy_comparison.py b/policyengine_api/services/economy_comparison.py index 2d5518827..4e8c5f573 100644 --- a/policyengine_api/services/economy_comparison.py +++ b/policyengine_api/services/economy_comparison.py @@ -13,6 +13,11 @@ def budgetary_impact(baseline: dict, reform: dict) -> dict: + """Return the reform's change in tax revenue, benefit spending and net budget. + + A positive ``budgetary_impact`` means the reform raises net revenue (tax + revenue change minus benefit spending change). + """ tax_revenue_impact = reform["total_tax"] - baseline["total_tax"] state_tax_revenue_impact = reform["total_state_tax"] - baseline["total_state_tax"] benefit_spending_impact = reform["total_benefits"] - baseline["total_benefits"] @@ -28,6 +33,11 @@ def budgetary_impact(baseline: dict, reform: dict) -> dict: def labor_supply_response(baseline: dict, reform: dict) -> dict: + """Return the reform's labour supply responses relative to the baseline. + + This covers the income and substitution effects in total, as shares of + original earnings, by income decile, and in weekly hours. + """ substitution_lsr = reform["substitution_lsr"] - baseline["substitution_lsr"] income_lsr = reform["income_lsr"] - baseline["income_lsr"] total_change = substitution_lsr + income_lsr @@ -103,6 +113,10 @@ def labor_supply_response(baseline: dict, reform: dict) -> dict: def detailed_budgetary_impact(baseline: dict, reform: dict, country_id: str) -> dict: + """Return baseline, reform and difference for each UK programme's budget. + + Only the UK reports programme-level totals; other countries get an empty dict. + """ result = {} if country_id == "uk": for program in baseline["programs"]: @@ -295,8 +309,9 @@ def poverty_impact(baseline: dict, reform: dict) -> dict: def compute_income_change( baseline_values: np.ndarray, reform_values: np.ndarray ) -> np.ndarray: - """Percentage income change with a floor of 1 on the baseline - to avoid division by zero for zero/negative incomes. + """Return each household's income change as a fraction of baseline income. + + The baseline is floored at 1 so zero or negative incomes do not divide by zero. """ absolute_change = reform_values - baseline_values capped_baseline = np.maximum(baseline_values, 1) @@ -304,6 +319,11 @@ def compute_income_change( def intra_decile_impact(baseline: dict, reform: dict) -> dict: + """Return the share of people in each income decile by size of income change. + + People fall into five buckets (gain or lose more or less than 5%, or no change + within 0.1%). ``all`` is the unweighted mean of the ten decile shares. + """ baseline_income = MicroSeries( baseline["household_net_income"], weights=baseline["household_weight"] ) @@ -356,6 +376,11 @@ def intra_decile_impact(baseline: dict, reform: dict) -> dict: def intra_wealth_decile_impact(baseline: dict, reform: dict) -> dict: + """Return the share of people in each wealth decile by size of income change. + + This uses the same buckets as ``intra_decile_impact``, grouped by household + wealth decile instead of income decile. + """ baseline_income = MicroSeries( baseline["household_net_income"], weights=baseline["household_weight"] ) @@ -508,6 +533,8 @@ def poverty_racial_breakdown(baseline: dict, reform: dict) -> dict: class UKConstituencyBreakdownByConstituency(BaseModel): + """One constituency's household income change and its map position (x, y).""" + average_household_income_change: float relative_household_income_change: float x: int @@ -515,11 +542,19 @@ class UKConstituencyBreakdownByConstituency(BaseModel): class UKConstituencyBreakdown(BaseModel): + """UK reform results by constituency. + + ``outcomes_by_region`` counts how many constituencies in the UK and each nation + fall into each income-change bucket. + """ + by_constituency: dict[str, UKConstituencyBreakdownByConstituency] outcomes_by_region: dict[str, dict[str, int]] class UKLocalAuthorityBreakdownByLA(BaseModel): + """One local authority's household income change and its map position (x, y).""" + average_household_income_change: float relative_household_income_change: float x: int @@ -527,6 +562,12 @@ class UKLocalAuthorityBreakdownByLA(BaseModel): class UKLocalAuthorityBreakdown(BaseModel): + """UK reform results by local authority. + + ``outcomes_by_region`` counts how many local authorities in the UK and each + nation fall into each income-change bucket. + """ + by_local_authority: dict[str, UKLocalAuthorityBreakdownByLA] outcomes_by_region: dict[str, dict[str, int]] @@ -534,6 +575,12 @@ class UKLocalAuthorityBreakdown(BaseModel): def uk_constituency_breakdown( baseline: dict, reform: dict, country_id: str, region: str | None = None ) -> UKConstituencyBreakdown | None: + """Return each UK constituency's average and relative household income change. + + Return None outside the UK or for a local-authority region. A constituency or + country region limits the output to that area. Constituency weights and + names come from PolicyEngine's Hugging Face datasets. + """ if country_id != "uk": return None @@ -663,6 +710,12 @@ def uk_constituency_breakdown( def uk_local_authority_breakdown( baseline: dict, reform: dict, country_id: str, region: str | None = None ) -> UKLocalAuthorityBreakdown | None: + """Return each UK local authority's average and relative household income change. + + Return None outside the UK or for a constituency region. A local-authority or + country region limits the output to that area. Local-authority weights and + names come from PolicyEngine's Hugging Face datasets. + """ if country_id != "uk": return None diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index fb7e81ecc..c2307b38d 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -1,3 +1,5 @@ +"""Economic impact service backed by cached reform impacts and simulation jobs.""" + import datetime import hashlib import json @@ -71,7 +73,7 @@ class ImpactAction(Enum): - """Enum for the action to take based on the status of an economic impact calculation.""" + """What to do with an economic impact request, given its stored calculation.""" COMPLETED = "completed" FAILED = "failed" @@ -98,6 +100,8 @@ class ImpactStatus(Enum): class SimulationOptions(BaseModel): + """The job payload sent to the Simulation entrypoint for one economy simulation.""" + spm: SPMSelection | None = None country: str scope: Literal["macro", "household"] = "macro" @@ -111,6 +115,12 @@ class SimulationOptions(BaseModel): class EconomicImpactSetupOptions(BaseModel): + """One resolved economic impact request. + + It holds the request's cache identity, the runtime bundle versions and the + submission claim token that guards against duplicate simulation jobs. + """ + submission_claim_id: str observability_id: str | None = None country_id: str @@ -130,8 +140,10 @@ class EconomicImpactSetupOptions(BaseModel): class EconomicImpactResult(BaseModel): - """Model for the result of an economic impact calculation. - Contains the status and the reform impact JSON. + """The status of an economic impact calculation and, once complete, its result. + + ``data`` holds the reform impact JSON when the status is ok, and ``message`` + explains an error. """ status: ImpactStatus @@ -181,6 +193,7 @@ class BudgetWindowEconomicImpactResult(BaseModel): model_config = {"frozen": True} def to_dict(self) -> dict[str, Any]: + """Return the response body; ``cache_status`` is internal and is left out.""" return { "status": self.status.value, "data": self.data, @@ -196,6 +209,7 @@ def to_dict(self) -> dict[str, Any]: def completed( cls, data: dict, *, cache_status: str | None = None ) -> "BudgetWindowEconomicImpactResult": + """Build a finished result with progress at 100.""" return cls( status=ImpactStatus.OK, data=data, @@ -214,6 +228,7 @@ def computing( message: str, cache_status: str | None = None, ) -> "BudgetWindowEconomicImpactResult": + """Build an in-progress result listing the done, running and queued years.""" return cls( status=ImpactStatus.COMPUTING, data=None, @@ -235,6 +250,7 @@ def failed( queued_years: list[str] | None = None, cache_status: str | None = None, ) -> "BudgetWindowEconomicImpactResult": + """Build an error result, copying the message into ``error`` and logging it.""" logger.log_struct({"message": message}, severity="ERROR") return cls( status=ImpactStatus.ERROR, @@ -249,9 +265,10 @@ def failed( class EconomyService: - """Service for calculating economic impact of policy reforms; this is connected - to the /economy route, which does not have its own table; therefore, it connects - with other services to access their respective tables + """Service that calculates the economic impact of policy reforms. + + It backs the /economy routes, which have no table of their own, so it reads + policies and cached reform impacts through the other services. """ def __init__( @@ -263,6 +280,7 @@ def __init__( budget_window_cache_: BudgetWindowCache | None = None, simulation_entrypoint_: SimulationEntrypointClient | None = None, ) -> None: + """Inject collaborators; omitted ones are built lazily or use the defaults.""" self._primary_session_factory = primary_session_factory self._injected_policy_service = policy_service_ self._injected_reform_impacts_service = reform_impacts_service_ @@ -332,10 +350,9 @@ def get_economic_impact( ) -> EconomicImpactResult: """Calculate the society-wide economic impact of a policy reform. - Returns a tuple of: - - status: "ok", "error", or "computing" - - reform_impact_json: The JSON object containing the economic impact data, or None if - the status is "computing" or "error". + Return a result whose status is "ok" with the impact JSON in ``data``, + "computing" while the calculation runs, or "error" with a ``message``. + Unexpected failures are logged and re-raised. """ set_runtime_context( country_id=country_id, @@ -404,6 +421,14 @@ def get_budget_window_economic_impact( target: Literal["general", "cliff"] = "general", max_active_years: int = BUDGET_WINDOW_MAX_ACTIVE_YEARS, ) -> BudgetWindowEconomicImpactResult: + """Return the multi-year budget-window impact, submitting a batch job if needed. + + The shared cache records each report's state (starting, submitted, completed or + failed) so concurrent requests share one batch job; later calls poll it and + report progress until it finishes. Raise CacheCoordinationError if the cache is + unavailable or submission ownership keeps changing, and SPMValidationError when + the SPM selection or a stored result cannot be certified. + """ set_runtime_context( country_id=country_id, policy_id=policy_id, @@ -1159,8 +1184,10 @@ def _get_most_recent_impact( self, setup_options: EconomicImpactSetupOptions, ) -> CachedReformImpact | None: - """Get the first impact from the reform_impacts table based on the provided setup options. - Returns the first impact if it exists, otherwise returns None. + """Return the newest cached impact for these options, or None if there is none. + + An impact with exactly this options hash wins; otherwise the newest impact + whose hash extends this one (with extra bundle fields) is used. """ previous_impacts = self._get_previous_impacts( country_id=setup_options.country_id, @@ -1208,7 +1235,7 @@ def _handle_execution_state( reform_impact: ReformImpact, execution: ModalSimulationExecution | None = None, ) -> EconomicImpactResult: - """Handle the state of the execution and return the appropriate status and result. + """Map an execution's state to the impact status and result to return. Supports both GCP Workflow statuses (SUCCEEDED, FAILED, ACTIVE) and Modal statuses (complete, failed, running, submitted). @@ -1710,7 +1737,7 @@ def _set_reform_impact_computing( setup_options: EconomicImpactSetupOptions, execution_id: str, ) -> None: - """In the reform_impact table, set the status of the impact to "computing".""" + """Record the impact as "computing" for this run; log and re-raise errors.""" try: self._reform_impacts.set_reform_impact( country_id=setup_options.country_id, @@ -1746,7 +1773,7 @@ def _set_reform_impact_complete( reform_impact_json: dict[str, Any], execution_id: str, ) -> None: - """In the reform_impact table, set the status of the impact to "ok" and store the reform impact JSON.""" + """Record the impact as "ok" with its result JSON; log and re-raise errors.""" try: self._reform_impacts.set_complete_reform_impact( country_id=setup_options.country_id, @@ -1778,7 +1805,7 @@ def _set_reform_impact_error( execution_id: str, error_code: str | None = None, ) -> None: - """In the reform_impact table, set the status of the impact to "error" and store the error message.""" + """Record the impact as "error" with its message; log and re-raise errors.""" try: self._reform_impacts.set_error_reform_impact( country_id=setup_options.country_id, diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py index 67e99ff5a..f8ab17a67 100644 --- a/policyengine_api/services/household_calculation_service.py +++ b/policyengine_api/services/household_calculation_service.py @@ -1,3 +1,5 @@ +"""Calculations for stored and request-supplied households, with result caching.""" + from __future__ import annotations import time @@ -39,6 +41,8 @@ @dataclass(frozen=True) class CalculationResult: + """One country-package calculation: the household, warnings and SPM receipt.""" + household: dict warnings: tuple[str, ...] = () spm_provenance: dict | None = None @@ -46,6 +50,8 @@ class CalculationResult: @dataclass(frozen=True) class HouseholdCalculationResult: + """A household calculation as returned to the route, and whether it was cached.""" + household: dict warnings: tuple[str, ...] = () cached: bool = False @@ -66,15 +72,18 @@ class PreparedHouseholdCalculation: class HouseholdNotFoundError(LookupError): - pass + """Raised when the stored household does not exist in the requested country.""" class PolicyNotFoundError(LookupError): - pass + """Raised when the stored policy does not exist in the requested country.""" class InvalidHouseholdInputsError(ValueError): + """Raised when a household or policy has inputs the country model does not know.""" + def __init__(self, invalid_inputs: list[Any]) -> None: + """Keep the unrecognised inputs so the route can report them.""" self.invalid_inputs = invalid_inputs super().__init__("Household or policy contains unrecognized inputs") @@ -95,6 +104,10 @@ def _validated_spm_provenance( def get_household_year(household: dict) -> int | str: + """Return the year of the first age entered for ``you``, or the current year. + + This is the period used when filling in missing yearly variables. + """ household_year: int | str = date.today().year household_age_list = list( household.get("people", {}).get("you", {}).get("age", {}).keys() @@ -109,6 +122,12 @@ def add_yearly_variables( country_id: str, countries: dict | None = None, ) -> dict: + """Add every yearly, monthly and eternal variable the household's entities lack. + + Input variables get their default value and computed variables get None, which + asks the model to calculate them, for the household's year. The household is + changed in place and returned. + """ if countries is None: from policyengine_api.country import COUNTRIES @@ -143,6 +162,7 @@ def __init__( cache: HouseholdCalculationCache | None = None, country_provider: Callable[[], dict] | None = None, ) -> None: + """Inject the session factory, cache and country source, or use the defaults.""" self._injected_primary_session_factory = primary_session_factory if cache is None: context = get_runtime_cache_context() @@ -225,6 +245,13 @@ def calculate_stored_household( *, on_accepted: Callable[[], object] | None = None, ) -> HouseholdCalculationResult: + """Calculate a saved household under a saved policy, reusing a cached result. + + Raise HouseholdNotFoundError or PolicyNotFoundError if either is missing in this + country, and InvalidHouseholdInputsError for unrecognised inputs. + ``on_accepted`` runs once the inputs have loaded, before the cache lookup and + calculation. + """ api_version = COUNTRY_PACKAGE_VERSIONS[country_id] with observability_runtime.span( HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_LOAD_INPUTS) diff --git a/policyengine_api/services/household_mirroring.py b/policyengine_api/services/household_mirroring.py index 940e9c4fa..bcf370718 100644 --- a/policyengine_api/services/household_mirroring.py +++ b/policyengine_api/services/household_mirroring.py @@ -36,7 +36,9 @@ class LegacyHouseholdMirror(Protocol): def mirror_legacy_household( self, snapshot: LegacyHouseholdSnapshot, - ) -> LegacyHouseholdPersistenceResult: ... + ) -> LegacyHouseholdPersistenceResult: + """Copy one committed v1 household into v2 and map its v1 id to the result.""" + ... class HouseholdMirrorUnavailableError(RuntimeError): diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index 6fdd7ee68..dd3b57010 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -1,3 +1,5 @@ +"""v1 household storage and the retained events that copy new households into v2.""" + from __future__ import annotations import copy @@ -68,6 +70,8 @@ class HouseholdMirrorEventIntegrityError(RuntimeError): class HouseholdPersistenceFailureCategory(StrEnum): + """Safe-to-report reasons a household write failed.""" + TIMEOUT = "timeout" UNAVAILABLE = "unavailable" INTEGRITY = "integrity" @@ -79,11 +83,13 @@ class HouseholdPersistenceError(RuntimeError): """Domain error containing only safe persistence-failure attributes.""" def __init__(self, category: HouseholdPersistenceFailureCategory) -> None: + """Record the failure category; the message never carries database details.""" super().__init__("Household persistence failed") self.category = HouseholdPersistenceFailureCategory(category) @property def retryable(self) -> bool: + """Whether a retry may succeed: true for timeouts and a database outage.""" return self.category in { HouseholdPersistenceFailureCategory.TIMEOUT, HouseholdPersistenceFailureCategory.UNAVAILABLE, @@ -91,6 +97,7 @@ def retryable(self) -> bool: @classmethod def from_exception(cls, error: Exception) -> HouseholdPersistenceError: + """Classify a database or unexpected error as a safe persistence error.""" if isinstance(error, SQLAlchemyTimeoutError): return cls(HouseholdPersistenceFailureCategory.TIMEOUT) if isinstance(error, OperationalError): @@ -104,6 +111,8 @@ def from_exception(cls, error: Exception) -> HouseholdPersistenceError: @dataclass(frozen=True) class PendingHouseholdMirrorEvent: + """A validated retained household copy event, and whether it was processed.""" + event_id: int snapshot: LegacyHouseholdSnapshot source_fingerprint_sha256: str @@ -112,6 +121,8 @@ class PendingHouseholdMirrorEvent: @dataclass(frozen=True) class HouseholdCreateResult: + """A newly stored v1 household, with its v2 copy snapshot and any event id.""" + household: Household snapshot: LegacyHouseholdSnapshot | None mirror_event_id: int | None @@ -124,6 +135,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -135,6 +147,10 @@ def get_household( country_id: str, household_id: int, ) -> Household | None: + """Return the stored household in this country, or None if there is none. + + Raise Exception if ``household_id`` is not a non-negative int. + """ if type(household_id) is not int or household_id < 0: raise Exception( f"Invalid household ID: {household_id}. Must be a positive integer." @@ -164,6 +180,11 @@ def create_household( record_mirror_event: bool = False, spm: object = None, ) -> HouseholdCreateResult: + """Store a v1 household, optionally recording an event to copy it into v2. + + Raise SPMValidationError for an invalid SPM selection before writing, and + HouseholdPersistenceError, carrying a safe category, if the write fails. + """ # Validate before the persistence boundary so typed input failures remain 400. household_json = household_storage_json(country_id, household_json, spm) try: diff --git a/policyengine_api/services/metadata_service.py b/policyengine_api/services/metadata_service.py index b71791a23..35bdf86b2 100644 --- a/policyengine_api/services/metadata_service.py +++ b/policyengine_api/services/metadata_service.py @@ -1,9 +1,14 @@ +"""Read access to the metadata document of each loaded country model.""" + from policyengine_api.country import COUNTRIES from policyengine_api.json_types import JSONObject class MetadataService: + """Service that returns a country's metadata from the loaded country models.""" + def get_metadata(self, country_id: str) -> JSONObject: + """Return the country's metadata; raise RuntimeError for an unknown country.""" country = COUNTRIES.get(country_id) if country is None: raise RuntimeError( diff --git a/policyengine_api/services/policy_mirroring.py b/policyengine_api/services/policy_mirroring.py index 217d2aba7..1de90ee88 100644 --- a/policyengine_api/services/policy_mirroring.py +++ b/policyengine_api/services/policy_mirroring.py @@ -36,7 +36,9 @@ class LegacyPolicyMirror(Protocol): def mirror_legacy_policy( self, snapshot: LegacyPolicySnapshot, - ) -> LegacyPolicyPersistenceResult: ... + ) -> LegacyPolicyPersistenceResult: + """Copy one committed v1 policy into v2 and map its v1 id to the result.""" + ... class PolicyMirrorUnavailableError(RuntimeError): diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index 714a59a38..d74f21c06 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -1,3 +1,5 @@ +"""v1 policy storage: lookup, label search and deduplicated creation.""" + from __future__ import annotations import copy @@ -43,6 +45,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -61,6 +64,11 @@ def get_policy( country_id: str, policy_id: int, ) -> Policy | None: + """Return the stored policy in this country, or None if there is none. + + Raise Exception if ``policy_id`` is not a non-negative int, and ValueError for + an empty ``country_id``. + """ self._validate_policy_id(policy_id) if not country_id: raise ValueError("country_id cannot be empty or None") @@ -85,6 +93,7 @@ def get_policy_json( country_id: str, policy_id: int, ) -> JSONValue: + """Return the policy's stored JSON, or None if the policy does not exist.""" policy = self.get_policy(country_id, policy_id) return None if policy is None else policy.policy_json @@ -113,6 +122,11 @@ def search_policies( *, unique_only: bool = False, ) -> list[Policy]: + """Return this country's policies whose label contains ``query``. + + With ``unique_only``, keep only the first policy for each label and content + hash. + """ with self._sessions() as session: results = list( session.scalars( @@ -142,6 +156,11 @@ def set_policy( *, prepare_for_mirroring: bool = False, ) -> PolicySetResult: + """Store a policy, or return the existing one with the same content and label. + + Raise ValueError for an unknown country. With ``prepare_for_mirroring`` the + result carries a snapshot of the stored row, for copying the policy into v2. + """ country_id = country_id.lower() if country_id not in COUNTRY_PACKAGE_VERSIONS: raise ValueError(f"Invalid country_id: {country_id}") diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 7a0c118f9..e199a18b0 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -21,6 +21,7 @@ class ReformImpactsService: """Preserve historical lookup contracts over an expiring cache repository.""" def __init__(self, cache: ReformImpactCache | None = None) -> None: + """Use this reform-impact cache, or one on the runtime cache if omitted.""" if cache is None: context = get_runtime_cache_context() cache = ReformImpactCache(context.client, context.namespace) @@ -30,6 +31,7 @@ def get_recent_reform_impacts( self, max_results: int, ) -> list[CachedReformImpact]: + """Return up to ``max_results`` cached impacts, most recently started first.""" return self._cache.recent(max_results) def get_all_reform_impacts( @@ -43,6 +45,7 @@ def get_all_reform_impacts( options_hash: str, api_version: str, ) -> list[CachedReformImpact]: + """Return cached impacts with exactly this options hash, newest first.""" return self._cache.matching( country_id=country_id, reform_policy_id=policy_id, @@ -66,6 +69,11 @@ def get_all_reform_impacts_by_options_hash_prefix( options_hash_prefix: str, api_version: str, ) -> list[CachedReformImpact]: + """Return cached impacts matching the options hash or the prefix pattern. + + ``options_hash_prefix`` is a SQL ``LIKE`` pattern. Exact hash matches come + first, then the most recently started. + """ return self._cache.matching( country_id=country_id, reform_policy_id=policy_id, @@ -184,6 +192,11 @@ def set_reform_impact( execution_id: str, observability_id: str | None = None, ) -> CachedReformImpact: + """Cache a newly submitted impact by execution ID so later requests can poll it. + + Raise ReformImpactHandoffError if the write fails, because the accepted job + could not otherwise be found again. + """ impact = CachedReformImpact( reform_impact_id=reform_impact_id(execution_id), country_id=country_id, @@ -219,6 +232,7 @@ def delete_reform_impact( time_period: str, options_hash: str, ) -> None: + """Remove this scope's in-flight (computing) impacts, ignoring cache errors.""" self._cache.delete_matching_computing( country_id=country_id, reform_policy_id=policy_id, @@ -242,6 +256,13 @@ def set_error_reform_impact( execution_id: str, error_code: str | None = None, ) -> CachedReformImpact | None: + """Mark this execution's cached impact as failed. + + Only ``execution_id`` selects the record; the scope arguments are kept for the + historical signature and ignored. An ``error_code`` records a typed, terminal + failure and clears the worker handle. Return the updated impact, or None if + it is missing or the write failed. + """ del ( country_id, policy_id, @@ -272,6 +293,11 @@ def set_complete_reform_impact( reform_impact_json: dict[str, Any], execution_id: str, ) -> CachedReformImpact | None: + """Store the finished result on this execution's cached impact. + + Only ``execution_id`` selects the record; the scope arguments are ignored. + Return the updated impact, or None if it is missing or the write failed. + """ del ( country_id, reform_policy_id, diff --git a/policyengine_api/services/report_output_alias_service.py b/policyengine_api/services/report_output_alias_service.py index 90d23c6bc..14feae677 100644 --- a/policyengine_api/services/report_output_alias_service.py +++ b/policyengine_api/services/report_output_alias_service.py @@ -1,3 +1,5 @@ +"""Aliases that redirect a legacy v1 report output ID to its canonical report output.""" + from sqlalchemy.orm import Session from policyengine_api.data.v1_models import LegacyReportOutputAlias, ReportOutput @@ -10,11 +12,17 @@ class ReportOutputAliasService: def get_alias( session: Session, legacy_report_output_id: int ) -> LegacyReportOutputAlias | None: + """Return the alias stored for this legacy report output ID, or None.""" return session.get(LegacyReportOutputAlias, legacy_report_output_id) def resolve_canonical_report_output_id( self, session: Session, requested_report_output_id: int ) -> int | None: + """Return the report output ID to serve for a requested ID, following any alias. + + Return None if there is neither an alias nor a report output with this ID, and + raise ValueError if the alias points to a missing report output. + """ alias = self.get_alias(session, requested_report_output_id) if alias is not None: canonical_id = alias.canonical_report_output_id @@ -32,6 +40,13 @@ def set_alias( legacy_report_output_id: int, canonical_report_output_id: int, ) -> bool: + """Point a legacy report output at a canonical one describing the same report. + + Repeating an existing alias is a no-op; the caller owns the transaction. Raise + ValueError if either output is missing, both IDs are the same, the legacy output + is already aliased elsewhere, or the two differ in country, simulations or year. + Always return True. + """ legacy = session.get(ReportOutput, legacy_report_output_id) if legacy is None: raise ValueError( diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index 0749bac83..96dd213b1 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -1,3 +1,5 @@ +"""v1 report outputs, with every write mirrored into the report's run history.""" + import json from dataclasses import dataclass from datetime import UTC, datetime @@ -85,6 +87,12 @@ def _canonicalize_society_wide_spm_output( @dataclass(frozen=True) class ReportOutputView: + """A report output as served: the record, the run to display and the ID to report. + + ``response_id`` is the requested ID when an outdated report was redirected to its + current equivalent; None means the record's own ID is reported. + """ + report_output: ReportOutput display_run: ReportOutputRun | None response_id: int | None = None @@ -92,6 +100,8 @@ class ReportOutputView: @dataclass(frozen=True) class ReportCreateResult: + """A report output view and whether this call created the report.""" + view: ReportOutputView created: bool @@ -103,6 +113,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory for reports and simulations, or the v1 one.""" self._injected_session_factory = session_factory self.report_spec_service = ReportSpecService() self.report_run_service = ReportRunService() @@ -118,6 +129,11 @@ def _utc_timestamp() -> datetime: @staticmethod def format_run_timestamp(value: datetime | str | None) -> str | None: + """Return a timestamp as an ISO 8601 UTC string ending in ``Z``, or None. + + Naive values are taken as UTC; an unparseable string is returned with ``Z`` + added. + """ if value is None: return None if isinstance(value, datetime): @@ -607,6 +623,10 @@ def create_or_reuse_report_output( simulation_2_id: int | None = None, year: str = "2025", ) -> ReportCreateResult: + """Return the current report for these simulations and year, creating it if new. + + Raise ValueError if a referenced simulation does not exist in this country. + """ with self._sessions.begin() as session: existing = self._find_existing_report_output( session, @@ -639,6 +659,13 @@ def resolve_report_output( country_id: str, report_output_id: int, ) -> ReportOutputView | None: + """Return the report output to serve for a requested ID, or None if missing. + + A report stored under an outdated cache version is served as the current report + for the same simulations and year (created if needed), with the requested ID + kept as ``response_id``. Raise Exception for an ID that is not a non-negative + int. + """ if type(report_output_id) is not int or report_output_id < 0: raise Exception( f"Invalid report output ID: {report_output_id}. " @@ -672,6 +699,12 @@ def update_report_output( output: dict | list | str | None = None, error_message: str | None = None, ) -> ReportOutputView | None: + """Apply the given status, output or error message and sync the run history. + + Return None when nothing was given. Raise LookupError if the report does not + exist, and ValueError when marking it running without an active run or + completing a US economy report whose SPM provenance is invalid. + """ values = { key: value for key, value in { diff --git a/policyengine_api/services/report_run_service.py b/policyengine_api/services/report_run_service.py index aaf6daf53..7db920901 100644 --- a/policyengine_api/services/report_run_service.py +++ b/policyengine_api/services/report_run_service.py @@ -1,3 +1,5 @@ +"""Run history for v1 report outputs, one numbered row per calculation attempt.""" + import uuid from datetime import UTC, datetime from typing import Any @@ -44,6 +46,12 @@ def create_report_output_run( version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, ) -> ReportOutputRun: + """Append the next numbered run to a report output and return it. + + The parent row is locked while the sequence number is chosen, and timestamps are + set to now according to ``status``. Raise ValueError if the report output does + not exist. The caller owns the transaction. + """ parent = session.scalar( select(ReportOutput) .where(ReportOutput.id == report_output_id) @@ -90,11 +98,13 @@ def create_report_output_run( def get_report_output_run( self, session: Session, run_id: str ) -> ReportOutputRun | None: + """Return the run with this ID, or None.""" return session.get(ReportOutputRun, run_id) def list_report_output_runs( self, session: Session, report_output_id: int ) -> list[ReportOutputRun]: + """Return all runs of a report output, oldest first.""" return list( session.scalars( select(ReportOutputRun) @@ -106,6 +116,7 @@ def list_report_output_runs( def get_newest_report_output_run( self, session: Session, report_output_id: int ) -> ReportOutputRun | None: + """Return the report output's highest-numbered run, or None if it has none.""" return session.scalar( select(ReportOutputRun) .where(ReportOutputRun.report_output_id == report_output_id) @@ -115,6 +126,12 @@ def get_newest_report_output_run( def select_display_run( self, session: Session, report_output: ReportOutput ) -> ReportOutputRun | None: + """Return the run whose result the report output should display, or None. + + Prefer the active run, then (for a failed report) the run holding its error, + then the latest successful run, then any run matching the report's result, and + finally the newest run. + """ runs = list( session.scalars( select(ReportOutputRun) diff --git a/policyengine_api/services/report_spec_service.py b/policyengine_api/services/report_spec_service.py index 6bcb82609..a3aa8444b 100644 --- a/policyengine_api/services/report_spec_service.py +++ b/policyengine_api/services/report_spec_service.py @@ -1,3 +1,5 @@ +"""Explicit specifications of what each v1 report output calculates.""" + import json from typing import Any, Literal @@ -14,12 +16,16 @@ class ReportSimulationInput(BaseModel): + """The population and policy of one simulation in a household report.""" + population_type: Literal["household", "geography"] population_id: str policy_id: int class HouseholdReportSpec(BaseModel): + """A household report: one household under one policy, or under two to compare.""" + country_id: str report_kind: Literal["household_single", "household_comparison"] time_period: str @@ -28,6 +34,11 @@ class HouseholdReportSpec(BaseModel): class EconomyReportSpec(BaseModel): + """An economy report: one region under a baseline and a reform policy. + + A single-policy report uses the same policy for both. + """ + country_id: str report_kind: Literal["economy_single", "economy_comparison"] time_period: str @@ -43,6 +54,8 @@ class EconomyReportSpec(BaseModel): class ReportSpecService: + """Service that builds, checks and stores report specs for v1 report outputs.""" + @staticmethod def _validate_schema_version(schema_version: int | None) -> None: if schema_version != REPORT_SPEC_SCHEMA_VERSION: @@ -120,6 +133,10 @@ def _validate_report_country( def infer_report_kind( simulation_1: Simulation, simulation_2: Simulation | None = None ) -> str: + """Return the report kind implied by the simulations' population type and count. + + Raise ValueError if the two simulations' population types differ or are unknown. + """ population_type = simulation_1.population_type if simulation_2 is not None and simulation_2.population_type != population_type: raise ValueError( @@ -140,6 +157,11 @@ def build_report_spec( target: Literal["general", "cliff"] = "general", options: dict[str, Any] | None = None, ) -> ReportSpec: + """Build the spec implied by a report output and its linked simulations. + + Raise ValueError if the simulations do not match the report's simulation IDs or + country, or if a comparison's simulations cover different households or regions. + """ self._validate_report_simulation_linkage( report_output, simulation_1, simulation_2 ) @@ -243,6 +265,11 @@ def _parse_report_spec(report_kind: str, raw_spec: dict) -> ReportSpec: def get_report_spec( self, session: Session, report_output_id: int ) -> ReportSpec | None: + """Return the stored spec for a report output, or None if it has none. + + Raise ValueError if the stored spec has an unsupported schema version or no + longer matches the report's linked simulations. + """ report_output = session.get(ReportOutput, report_output_id) if report_output is None or report_output.report_spec_json is None: return None @@ -264,6 +291,11 @@ def set_report_spec( report_spec_status: Literal["explicit", "backfilled_assumed"], schema_version: int = REPORT_SPEC_SCHEMA_VERSION, ) -> bool: + """Store a report spec after checking it matches the linked simulations. + + Raise ValueError for an unknown status or schema version, a missing report, or a + mismatched spec. Always return True; the caller owns the transaction. + """ if report_spec_status not in REPORT_SPEC_STATUSES: raise ValueError(f"Unsupported report spec status: {report_spec_status}") self._validate_schema_version(schema_version) diff --git a/policyengine_api/services/simulation_run_service.py b/policyengine_api/services/simulation_run_service.py index 677f2b6d5..91b3929fa 100644 --- a/policyengine_api/services/simulation_run_service.py +++ b/policyengine_api/services/simulation_run_service.py @@ -1,3 +1,5 @@ +"""Run history for v1 simulations, one numbered row per calculation attempt.""" + import uuid from typing import Any @@ -16,6 +18,8 @@ class SimulationRunService: + """Simulation-run operations performed through a caller-owned ORM session.""" + def create_simulation_run( self, session: Session, @@ -31,6 +35,12 @@ def create_simulation_run( version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, ) -> SimulationRun: + """Append the next numbered run to a simulation and return it. + + The parent row is locked while the sequence number is chosen; timestamps start + unset. Raise ValueError if the simulation does not exist. The caller owns the + transaction. + """ parent = session.scalar( select(Simulation).where(Simulation.id == simulation_id).with_for_update() ) @@ -78,6 +88,7 @@ def get_simulation_run( session: Session, run_id: str, ) -> SimulationRun | None: + """Return the run with this ID, or None.""" return session.get(SimulationRun, run_id) def list_simulation_runs( @@ -85,6 +96,7 @@ def list_simulation_runs( session: Session, simulation_id: int, ) -> list[SimulationRun]: + """Return all runs of a simulation, oldest first.""" return list( session.scalars( select(SimulationRun) @@ -98,6 +110,7 @@ def get_newest_simulation_run( session: Session, simulation_id: int, ) -> SimulationRun | None: + """Return the simulation's highest-numbered run, or None if it has none.""" return session.scalar( select(SimulationRun) .where(SimulationRun.simulation_id == simulation_id) @@ -109,6 +122,10 @@ def select_display_run( session: Session, simulation: Simulation, ) -> SimulationRun | None: + """Return the run whose result the simulation should display, or None. + + Prefer the active run, then the latest successful run, then the newest run. + """ if simulation.active_run_id: active_run = self.get_simulation_run(session, simulation.active_run_id) if active_run is not None: diff --git a/policyengine_api/services/simulation_service.py b/policyengine_api/services/simulation_service.py index 3ef805673..da9e0a8b0 100644 --- a/policyengine_api/services/simulation_service.py +++ b/policyengine_api/services/simulation_service.py @@ -1,3 +1,5 @@ +"""v1 simulations, with every write mirrored into the simulation's run history.""" + import json import uuid from dataclasses import dataclass @@ -17,6 +19,8 @@ @dataclass(frozen=True) class SimulationCreateResult: + """A simulation and whether this call created it.""" + simulation: Simulation created: bool @@ -28,6 +32,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -207,6 +212,12 @@ def get_or_create_simulation( population_type: str, policy_id: int, ) -> SimulationCreateResult: + """Return the simulation for this population and policy, creating it if needed. + + A numeric household ID also matches its stored aliases. Stored output that + predates the current SPM receipt format is cleared and the simulation reset to + pending, so it is calculated again. + """ numeric_identity = canonical_numeric_household_id( country_id, population_id, population_type ) @@ -252,6 +263,10 @@ def get_simulation( country_id: str, simulation_id: int, ) -> Simulation | None: + """Return the simulation in this country, or None. + + Raise Exception if ``simulation_id`` is not a non-negative int. + """ if type(simulation_id) is not int or simulation_id < 0: raise Exception( f"Invalid simulation ID: {simulation_id}. Must be a positive integer." @@ -267,6 +282,12 @@ def update_simulation( output: dict | list | str | None = None, error_message: str | None = None, ) -> Simulation | None: + """Apply the given status, output or error message and sync the run history. + + Return None when nothing was given. The simulation is restamped with the current + country package version. Raise ValueError for output with legacy ``spm_config`` + or an invalid SPM receipt, and LookupError if the simulation does not exist. + """ values = { key: value for key, value in { diff --git a/policyengine_api/services/simulation_spec_service.py b/policyengine_api/services/simulation_spec_service.py index f6c5a7373..d75f6a76f 100644 --- a/policyengine_api/services/simulation_spec_service.py +++ b/policyengine_api/services/simulation_spec_service.py @@ -1,3 +1,5 @@ +"""Explicit specifications of what each v1 simulation calculates.""" + import json from typing import Literal @@ -10,6 +12,8 @@ class SimulationSpec(BaseModel): + """The population and policy one simulation runs, as stored on its row.""" + country_id: str population_id: str population_type: Literal["household", "geography"] @@ -17,6 +21,8 @@ class SimulationSpec(BaseModel): class SimulationSpecService: + """Service that builds, checks and stores the spec saved on a v1 simulation.""" + def _validate_schema_version(self, schema_version: int | None) -> None: if schema_version != SIMULATION_SPEC_SCHEMA_VERSION: raise ValueError( @@ -38,6 +44,7 @@ def _validate_simulation_spec_matches_model( raise ValueError("Simulation spec must match the linked simulation row") def build_simulation_spec(self, simulation: Simulation) -> SimulationSpec: + """Return the spec implied by a simulation's country, population and policy.""" return SimulationSpec( country_id=simulation.country_id, population_id=simulation.population_id, @@ -50,6 +57,11 @@ def get_simulation_spec( session: Session, simulation_id: int, ) -> SimulationSpec | None: + """Return the stored spec for a simulation, or None if it has none. + + Raise ValueError if the stored spec has an unsupported schema version or no + longer matches the simulation row. + """ simulation = session.get(Simulation, simulation_id) if simulation is None or simulation.simulation_spec_json is None: return None @@ -71,6 +83,11 @@ def set_simulation_spec( simulation_spec: SimulationSpec, schema_version: int = SIMULATION_SPEC_SCHEMA_VERSION, ) -> bool: + """Store a spec on a simulation after checking it matches the simulation row. + + Raise ValueError for an unsupported schema version, a missing simulation or a + mismatched spec. Always return True; the caller owns the transaction. + """ self._validate_schema_version(schema_version) simulation = session.get(Simulation, simulation_id) if simulation is None: diff --git a/policyengine_api/services/user_policy_mirroring.py b/policyengine_api/services/user_policy_mirroring.py index dddadf499..4b28a7aac 100644 --- a/policyengine_api/services/user_policy_mirroring.py +++ b/policyengine_api/services/user_policy_mirroring.py @@ -40,6 +40,8 @@ class LegacyUserPolicyMirror(Protocol): + """Supabase transaction service used after a v1 saved-policy commit.""" + def mirror_legacy_user_policy( self, snapshot: LegacyUserPolicySnapshot, @@ -47,7 +49,9 @@ def mirror_legacy_user_policy( *, source_revision: int, changed_fields: frozenset[str], - ) -> LegacyUserPolicyPersistenceResult: ... + ) -> LegacyUserPolicyPersistenceResult: + """Copy one v1 saved-policy revision, with its reform policy, into v2.""" + ... class UserPolicyMirrorUnavailableError(RuntimeError): @@ -216,7 +220,12 @@ def mirror_pending_user_policy_events_after_commit( | None = None, mirror_factory: Callable[[], LegacyUserPolicyMirror] | None = None, ) -> LegacyUserPolicyPersistenceResult: - """Synchronously apply retained source events through one request's revision.""" + """Copy this saved policy's retained events into v2, up to the request's revision. + + Runs synchronously after the v1 commit and returns the result for + ``through_revision``. Raise UserPolicyMirrorUnavailableError, after logging the + failure, if an event cannot be copied. + """ selected_event_service = event_service or UserPolicyService() selected_reform_loader = ( reform_snapshot_loader or PolicyService().get_policy_snapshot diff --git a/policyengine_api/services/user_policy_service.py b/policyengine_api/services/user_policy_service.py index 99ffc46ea..98d26eb12 100644 --- a/policyengine_api/services/user_policy_service.py +++ b/policyengine_api/services/user_policy_service.py @@ -1,3 +1,5 @@ +"""v1 saved policies and the retained revision events that copy them into v2.""" + from __future__ import annotations from dataclasses import dataclass @@ -64,6 +66,8 @@ class UserPolicyMirrorEventIntegrityError(RuntimeError): class UserPolicyPersistenceFailureCategory(StrEnum): + """Safe-to-report reasons a saved-policy write failed.""" + TIMEOUT = "timeout" UNAVAILABLE = "unavailable" INTEGRITY = "integrity" @@ -75,11 +79,13 @@ class UserPolicyPersistenceError(RuntimeError): """Domain error containing only safe persistence-failure attributes.""" def __init__(self, category: UserPolicyPersistenceFailureCategory) -> None: + """Record the failure category; the message never carries database details.""" super().__init__("Saved-policy persistence failed") self.category = UserPolicyPersistenceFailureCategory(category) @property def retryable(self) -> bool: + """Whether a retry may succeed: true for timeouts and a database outage.""" return self.category in { UserPolicyPersistenceFailureCategory.TIMEOUT, UserPolicyPersistenceFailureCategory.UNAVAILABLE, @@ -87,6 +93,7 @@ def retryable(self) -> bool: @classmethod def from_exception(cls, error: Exception) -> UserPolicyPersistenceError: + """Classify a database or unexpected error as a safe persistence error.""" if isinstance(error, SQLAlchemyTimeoutError): return cls(UserPolicyPersistenceFailureCategory.TIMEOUT) if isinstance(error, OperationalError): @@ -100,6 +107,8 @@ def from_exception(cls, error: Exception) -> UserPolicyPersistenceError: @dataclass(frozen=True) class PendingUserPolicyMirrorEvent: + """A validated retained saved-policy event: one revision and what it changed.""" + event_id: int source_revision: int event_type: str @@ -110,6 +119,8 @@ class PendingUserPolicyMirrorEvent: @dataclass(frozen=True) class UserPolicyCreateResult: + """A saved policy, whether this call created it, and any mirror revision.""" + user_policy: UserPolicy created: bool snapshot: LegacyUserPolicySnapshot | None @@ -118,6 +129,8 @@ class UserPolicyCreateResult: @dataclass(frozen=True) class UserPolicyUpdateResult: + """An updated saved policy, the fields set, and its mirror revision if recorded.""" + user_policy: UserPolicy snapshot: LegacyUserPolicySnapshot | None changed_fields: frozenset[str] @@ -131,6 +144,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -277,6 +291,12 @@ def create_or_get_user_policy( *, record_mirror_event: bool = False, ) -> UserPolicyCreateResult: + """Return the saved policy with these identity fields, creating it if missing. + + With ``record_mirror_event`` the row is locked and a create revision event is + recorded, even for an existing policy. Raise UserPolicyPersistenceError, + carrying a safe category, if the write fails. + """ try: with self._sessions.begin() as session: user_policy = self._find_matching_user_policy( @@ -315,6 +335,7 @@ def list_user_policies( country_id: str, user_id: str, ) -> list[UserPolicy]: + """Return every saved policy this user has in this country.""" with self._sessions() as session: return list( session.scalars( @@ -333,6 +354,12 @@ def update_user_policy( *, record_mirror_event: bool = False, ) -> UserPolicyUpdateResult | None: + """Set the given fields on a saved policy; return None if not in this country. + + With ``record_mirror_event`` the row is locked and an update revision event is + recorded. Raise UserPolicyPersistenceError, carrying a safe category, if the + write fails. + """ try: with self._sessions.begin() as session: statement = select(UserPolicy).where( diff --git a/policyengine_api/services/user_service.py b/policyengine_api/services/user_service.py index f27fa7d1e..e60a1a11b 100644 --- a/policyengine_api/services/user_service.py +++ b/policyengine_api/services/user_service.py @@ -1,3 +1,5 @@ +"""v1 user profiles, keyed by internal user ID and by Auth0 ID.""" + from __future__ import annotations from typing import TYPE_CHECKING @@ -18,6 +20,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -31,6 +34,10 @@ def create_profile( username: str | None, user_since: int, ) -> tuple[bool, UserProfile]: + """Create a profile for this Auth0 ID, or return the one that already exists. + + Return ``(created, profile)``. + """ with self._sessions.begin() as session: return self._create_profile( session, @@ -67,6 +74,10 @@ def get_profile( auth0_id: str | None = None, user_id: int | str | None = None, ) -> UserProfile | None: + """Return the profile with this user ID, or else this Auth0 ID, or None. + + ``user_id`` wins when both are given. Raise ValueError if neither is given. + """ if auth0_id is None and user_id is None: raise ValueError("you must specify either auth0_id or user_id") with self._sessions() as session: @@ -92,6 +103,10 @@ def update_profile( username: str | None, user_since: int | None, ) -> UserProfile | None: + """Set the given non-None fields on a profile; return None if it does not exist. + + Raise ValueError if ``user_id`` is None. + """ if user_id is None: raise ValueError("you must specify either auth0_id or user_id") with self._sessions.begin() as session: From 8e6031308fb0f72963477ff43cc5ffd9bda45425 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:23:59 +0400 Subject: [PATCH 10/12] Enable ruff group 4 (docstrings) Select the pydocstyle rules (D), keeping D211 and D212 of the two mutually exclusive pairs. Test files stay exempt from the missing-docstring rules (D1) but follow the docstring style rules. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 499724a2c..6fa19e213 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,7 +91,12 @@ select = [ "TC", # Group 3: annotations. "ANN", + # Group 4: docstrings. + "D", ] +# Mutually exclusive pydocstyle pairs: keep D211 (no blank line before a class +# docstring) and D212 (summary on the first line). +ignore = ["D203", "D213"] [tool.ruff.lint.per-file-ignores] # Alembic revisions must stay exactly as Alembic generates them From 50ab7d3583941724a8f8e72b4904d51c6656b3e7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:26:17 +0400 Subject: [PATCH 11/12] Document the validator proof fixtures and drop a redundant return The two validator fixtures were the last D103 findings. The disabled cache backend's get() now ends at its docstring instead of suppressing RET501 on an explicit return None. Co-Authored-By: Claude Opus 5.5 --- policyengine_api/runtime_cache/fake.py | 1 - .../api.v2.validators-transformations-no-db/bad_validators.py | 1 + .../api.v2.validators-transformations-no-db/good_validators.py | 1 + 3 files changed, 2 insertions(+), 1 deletion(-) diff --git a/policyengine_api/runtime_cache/fake.py b/policyengine_api/runtime_cache/fake.py index f7c50f834..89b83a0c7 100644 --- a/policyengine_api/runtime_cache/fake.py +++ b/policyengine_api/runtime_cache/fake.py @@ -212,7 +212,6 @@ class DisabledCacheBackend(InMemoryCacheBackend): def get(self, key: str) -> None: """Return None: every read is a miss.""" - return None # noqa: RET501 # explicit None mirrors the Redis get contract def set(self, key: str, value: str, **kwargs: object) -> bool: """Discard the value and report the write as not stored.""" diff --git a/rules/fixtures/api.v2.validators-transformations-no-db/bad_validators.py b/rules/fixtures/api.v2.validators-transformations-no-db/bad_validators.py index a093c9b3c..9680bbdcd 100644 --- a/rules/fixtures/api.v2.validators-transformations-no-db/bad_validators.py +++ b/rules/fixtures/api.v2.validators-transformations-no-db/bad_validators.py @@ -6,6 +6,7 @@ def validate_policy_exists(session: Session, policy_id: int) -> Policy: + """Load the policy and fail if it doesn't exist.""" policy = session.exec(select(Policy).where(Policy.id == policy_id)).first() if policy is None: raise ValueError(policy_id) diff --git a/rules/fixtures/api.v2.validators-transformations-no-db/good_validators.py b/rules/fixtures/api.v2.validators-transformations-no-db/good_validators.py index 5b04ea204..aae3fa903 100644 --- a/rules/fixtures/api.v2.validators-transformations-no-db/good_validators.py +++ b/rules/fixtures/api.v2.validators-transformations-no-db/good_validators.py @@ -4,6 +4,7 @@ def validate_policy_exists(policy: Policy | None, policy_id: int) -> Policy: + """Fail if the already-loaded policy is missing.""" if policy is None: raise ValueError(policy_id) return policy From 90d5fe96ad8c404dc3442eb833e1f34318061fe7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:26:17 +0400 Subject: [PATCH 12/12] Register D in the ruff-lint rule api.quality.ruff-lint now lists D, and the rendered rules page and AGENTS.md Rules block are regenerated for the new tool hash. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- changelog.d/ruff-docstrings.changed.md | 1 + docs/generated/rules.md | 8 ++++---- rules/registry.toml | 2 +- 4 files changed, 7 insertions(+), 6 deletions(-) create mode 100644 changelog.d/ruff-docstrings.changed.md diff --git a/AGENTS.md b/AGENTS.md index 543b529fc..bd7898037 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -122,7 +122,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. - `api.quality.environment-variables-load`: Required environment variables and credentials load in a deployed-like environment. (blocking: CI job `test_env_vars`) - `api.quality.mypy-v2`: Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. (blocking: mypy) - `api.quality.ruff-format`: All Python files are formatted with ruff format. (blocking: ruff format) -- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)) +- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)) - `api.quality.runtime-evaluated-models-resolve`: Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. (blocking: pytest `tests/unit/test_runtime_evaluated_annotations.py`) - `api.quality.tests-pass`: The unit, contract and to-refactor test suites pass. (blocking: CI job `test`) - `api.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**`) diff --git a/changelog.d/ruff-docstrings.changed.md b/changelog.d/ruff-docstrings.changed.md new file mode 100644 index 000000000..cbdc5aa29 --- /dev/null +++ b/changelog.d/ruff-docstrings.changed.md @@ -0,0 +1 @@ +Enabled ruff's docstring rules (D) and documented every module, class and function outside tests; published models and routes now carry these descriptions in their JSON schema and OpenAPI, with no structural schema change. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 4778b566a..3ceb6b20c 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -269,7 +269,7 @@ flowchart LR | `api.quality.environment-variables-load` | Required environment variables and credentials load in a deployed-like environment. | mechanical | enforced | yes | CI job `test_env_vars` | `pr.yml:test_env_vars` | — | | `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — | | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — | -| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — | +| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — | | `api.quality.runtime-evaluated-models-resolve` | Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** | | `api.quality.tests-pass` | The unit, contract and to-refactor test suites pass. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — | | `api.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**` | `pr.yml:lint` | — | @@ -289,7 +289,7 @@ flowchart LR e_ruff_format["ruff format"] j_CI_pr_yml_lint["CI pr.yml:lint"] r_api_quality_ruff_lint("api.quality.ruff-lint (enforced)") - e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"] + e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"] r_api_quality_runtime_evaluated_models_resolve("api.quality.runtime-evaluated-models-resolve (enforced)") e_pytest_tests_unit_test_runtime_evaluated_annotations_py["pytest tests/unit/test_runtime_evaluated_annotations.py"] j_CI_pr_yml_test["CI pr.yml:test"] @@ -305,8 +305,8 @@ flowchart LR e_mypy --> j_CI_pr_yml_typecheck_v2 r_api_quality_ruff_format --> e_ruff_format e_ruff_format --> j_CI_pr_yml_lint - r_api_quality_ruff_lint --> e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ - e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint + r_api_quality_ruff_lint --> e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ + e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint r_api_quality_runtime_evaluated_models_resolve --> e_pytest_tests_unit_test_runtime_evaluated_annotations_py e_pytest_tests_unit_test_runtime_evaluated_annotations_py --> j_CI_pr_yml_test r_api_quality_tests_pass --> e_CI_job_test diff --git a/rules/registry.toml b/rules/registry.toml index 08096ff8e..7bfb2f027 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -679,7 +679,7 @@ autofix = true [rule.enforcer] kind = "ruff" mode = "lint" -codes = ["ANN", "E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"] +codes = ["ANN", "D", "E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"] [[rule.ci]] workflow = ".github/workflows/pr.yml" job = "lint"