Refactor the economy and calculation services within ruff's complexity limits

changelog.d/ruff-complexity-economy.changed.md

@@ -0,0 +1 @@
+Refactored the economy comparison and service, the reform-impact and budget-window caches, the SPM and worker-SPM modules and the household calculation service under ruff's complexity rules (C901, PLR): long functions are split into named steps, related arguments travel as small value objects, and magic numbers are named, with characterization tests pinning the behaviour first; responses, cache keys, stored records and published schemas are unchanged.

policyengine_api/routes/economy_routes.py

@@ -14,6 +14,7 @@
     parse_multidict_query,
 )
 from policyengine_api.response_factory import _make_error_response
+from policyengine_api.services.economic_impact_request import EconomicImpactRequest
 from policyengine_api.services.economy_service import (
     BudgetWindowEconomicImpactResult,
     EconomicImpactResult,
@@ -90,15 +91,18 @@ def get_economic_impact(
         query = parse_multidict_query(AnnualEconomyQuery, request.args)
         economic_impact_result: EconomicImpactResult = (
             economy_service.get_economic_impact(
-                country_id=country_id,
-                policy_id=policy_id,
-                baseline_policy_id=baseline_policy_id,
-                region=query.region,
-                dataset="default",
+                impact_request=EconomicImpactRequest(
+                    country_id=country_id,
+                    policy_id=policy_id,
+                    baseline_policy_id=baseline_policy_id,
+                    region=query.region,
+                    dataset="default",
+                    options=query.calculation_options(),
+                    api_version=query.version
+                    or COUNTRY_PACKAGE_VERSIONS.get(country_id),
+                    target=query.target,
+                ),
                 time_period=query.time_period,
-                options=query.calculation_options(),
-                api_version=query.version or COUNTRY_PACKAGE_VERSIONS.get(country_id),
-                target=query.target,
             )
         )
     except ValueError as error:
@@ -146,16 +150,19 @@ def get_budget_window_economic_impact(
         query = parse_multidict_query(BudgetWindowEconomyQuery, request.args)
         economic_impact_result: BudgetWindowEconomicImpactResult = (
             economy_service.get_budget_window_economic_impact(
-                country_id=country_id,
-                policy_id=policy_id,
-                baseline_policy_id=baseline_policy_id,
-                region=query.region,
-                dataset="default",
+                impact_request=EconomicImpactRequest(
+                    country_id=country_id,
+                    policy_id=policy_id,
+                    baseline_policy_id=baseline_policy_id,
+                    region=query.region,
+                    dataset="default",
+                    options=query.calculation_options(),
+                    api_version=query.version
+                    or COUNTRY_PACKAGE_VERSIONS.get(country_id),
+                    target=query.target,
+                ),
                 start_year=query.start_year,
                 window_size=query.window_size,
-                options=query.calculation_options(),
-                api_version=query.version or COUNTRY_PACKAGE_VERSIONS.get(country_id),
-                target=query.target,
             )
         )
     except ValueError as error:

policyengine_api/routes/household_routes.py

@@ -26,6 +26,7 @@
 )
 from policyengine_api.response_factory import _make_error_response
 from policyengine_api.services.household_calculation_service import (
+    HouseholdCalculationOptions,
     HouseholdCalculationResult,
     HouseholdCalculationService,
     HouseholdNotFoundError,
@@ -214,14 +215,16 @@ def wrapped(
                         country_id,
                         payload.get("household", {}),
                         payload.get("policy", {}),
-                        add_missing=add_missing,
-                        **(
-                            {
-                                "spm": g.spm,
-                                "spm_requested": g.spm_requested,
-                            }
-                            if g.spm is not None
-                            else {}
+                        options=HouseholdCalculationOptions(
+                            add_missing=add_missing,
+                            **(
+                                {
+                                    "spm": g.spm,
+                                    "spm_requested": g.spm_requested,
+                                }
+                                if g.spm is not None
+                                else {}
+                            ),
                         ),
                     )
                 )

policyengine_api/runtime_cache/fake.py

@@ -8,6 +8,10 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+# The claim store's compare-and-delete script takes one key and one argument:
+# the token the key must still hold.
+COMPARE_AND_DELETE_ARGUMENT_COUNT = 2
+
 
 class InMemoryCacheBackend:
     """Small Redis-like fake with explicit time advancement and atomic pipelines."""
@@ -153,7 +157,7 @@ def eval(
         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:
+        if numkeys != 1 or len(keys_and_args) != COMPARE_AND_DELETE_ARGUMENT_COUNT:
             msg = "fake supports one-key compare-and-delete only"
             raise ValueError(msg)
         key, expected = keys_and_args

policyengine_api/runtime_cache/household_calculations.py

@@ -1,7 +1,7 @@
 """Recoverable calculated-household result caching."""
 
 from dataclasses import asdict, dataclass
-from typing import Any
+from typing import Any, TypeGuard
 
 from pydantic import ValidationError
 
@@ -74,36 +74,19 @@ def get(
         different SPM configuration than ``identity`` asks for is also a miss.
         """
         payload = self._cache.get(asdict(identity))
-        if not isinstance(payload, dict):
-            return None
-        household = payload.get("household")
-        warnings = payload.get("warnings")
-        if not isinstance(household, dict) or not isinstance(warnings, list):
-            return None
-        if not all(isinstance(warning, str) for warning in warnings):
-            return None
-        if "spm_config" in payload:
+        if not _is_well_formed_entry(payload):
             return None
         spm_provenance = payload.get("spm_provenance")
-        if spm_provenance is not None and not isinstance(spm_provenance, dict):
-            return None
         if identity.spm is None:
             if spm_provenance is not None:
                 return None
         else:
-            try:
-                receipt = validate_spm_calculation_provenance(spm_provenance)
-            except (ValidationError, ValueError):
-                return None
-            if (
-                resolved_spm_configuration(receipt).model_dump(mode="json")
-                != identity.spm
-            ):
+            spm_provenance = _certified_receipt(spm_provenance, identity.spm)
+            if spm_provenance is None:
                 return None
-            spm_provenance = receipt.model_dump(mode="json", by_alias=True)
         return CachedHouseholdCalculation(
-            household=household,
-            warnings=tuple(warnings),
+            household=payload["household"],
+            warnings=tuple(payload["warnings"]),
             spm_provenance=spm_provenance,
         )
 
@@ -130,3 +113,36 @@ def set(
         payload = asdict(value)
         payload["spm_provenance"] = receipt.model_dump(mode="json", by_alias=True)
         return self._cache.set(asdict(identity), payload)
+
+
+def _is_well_formed_entry(payload: object) -> TypeGuard[dict[str, Any]]:
+    """Return whether a cached payload has the current calculated-household shape.
+
+    That is a household object, a list of string warnings, no legacy
+    ``spm_config`` and an SPM receipt that is an object when present.
+    """
+    if not isinstance(payload, dict):
+        return False
+    household = payload.get("household")
+    warnings = payload.get("warnings")
+    if not isinstance(household, dict) or not isinstance(warnings, list):
+        return False
+    if not all(isinstance(warning, str) for warning in warnings):
+        return False
+    if "spm_config" in payload:
+        return False
+    spm_provenance = payload.get("spm_provenance")
+    return spm_provenance is None or isinstance(spm_provenance, dict)
+
+
+def _certified_receipt(
+    spm_provenance: dict[str, Any] | None, expected_spm: dict[str, Any]
+) -> dict[str, Any] | None:
+    """Return the receipt in canonical form, or None unless it certifies the SPM."""
+    try:
+        receipt = validate_spm_calculation_provenance(spm_provenance)
+    except (ValidationError, ValueError):
+        return None
+    if resolved_spm_configuration(receipt).model_dump(mode="json") != expected_spm:
+        return None
+    return receipt.model_dump(mode="json", by_alias=True)