From 3c0a9493c62499f6b317829f157657f99d34a98a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:11:23 +0400 Subject: [PATCH 1/2] Open the simulation runtime's spans with the typed span scope #879 already opens the annual run's eight request, dataset, policy and simulation-build spans in simulation_runtime.py with scopes.span, and annotates _build_simulation as returning the StageReportingSimulation that _report_calculation_start needs. This commit opens the ninth span, credential setup in setup_gcp_credentials, the same way, and documents why _build_simulation returns a StageReportingSimulation in both branches (ArtifactBaselineSimulation subclasses it). Neither change alters what a run does, and the executor's baseline stays at 77 entries. Co-Authored-By: Claude Opus 5.5 --- .../policyengine_simulation_executor/simulation_runtime.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index f21114f4b..e1bb79e45 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -133,7 +133,7 @@ def setup_gcp_credentials( with contextlib.ExitStack() as credentials_file_stack: try: credential_span = ( - runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) + scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) if runtime is not None else contextlib.nullcontext() ) @@ -789,6 +789,11 @@ def _build_simulation( policy: dict[str, Any] | None, region: RegionResolution, ) -> StageReportingSimulation: + """Build a simulation that reports when its calculation starts. + + A run with a deterministic baseline id gets an ArtifactBaselineSimulation, + which is a StageReportingSimulation that can load a precomputed baseline. + """ from policyengine_simulation_executor.baseline_artifacts import ( ArtifactBaselineSimulation, BaselineIdentityRequest, From fde2bcbee1acdab46f04e80266f8c023c18bf4be Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 19:11:07 +0400 Subject: [PATCH 2/2] Type the simulation runtime's policy, region and dataset reads The last 5 Pyright errors in simulation_runtime.py, without changing what a run does: - _compile_interval_policy: compile_reform returns None only for an empty policy, which the function has already returned, so a None result now returns the policy unchanged instead of failing on .items(). - _build_uk_weight_replacement_region checks the region prefix against a tuple instead of a set display. A set display is typed set[str], so the check didn't narrow prefix to the two UK region types Region accepts; the tuple keeps their literal types. - The UK weight-matrix check counts households from the input dataset's entity_data, which Dataset types as an optional BaseModel. It now reads entity_data with getattr and requires a dict, raising TypeError with the dataset's name when none is loaded; a loaded dataset counts the same households. - _build_simulation reassigned its `policy: dict | None` parameter to the compiled policy, which may be a Policy; the result now has its own name, runtime_policy, used where the reassigned name was. The executor's baseline drops from 77 entries to 72. Co-Authored-By: Claude Opus 5.5 --- docs/generated/rules.md | 2 +- .../simulation_runtime.py | 21 +++++++++++++------ rules/baselines/pyright-executor.txt | 5 ----- rules/registry.toml | 2 +- 4 files changed, 17 insertions(+), 13 deletions(-) diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 686c1ea01..8e880c3c8 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt` -- `sim.quality.pyright-executor`: 77 remaining in `rules/baselines/pyright-executor.txt` +- `sim.quality.pyright-executor`: 72 remaining in `rules/baselines/pyright-executor.txt` - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt` ### Exceptions and review dates diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index e1bb79e45..2336e48f1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -290,6 +290,9 @@ def _compile_interval_policy( ) compiled = compile_reform(policy, year=year, model_version=country_module.model) + if compiled is None: + # compile_reform returns None only for an empty policy, returned above. + return policy parameter_values = [] for path, periods in compiled.items(): dated_values = [ @@ -363,7 +366,9 @@ def _build_uk_weight_replacement_region(region_code: str) -> Region | None: return None prefix, value = region_code.split("/", maxsplit=1) - if prefix not in {"constituency", "local_authority"}: + # A tuple, unlike a set display, keeps the literal types, so the check + # narrows prefix to the two UK region types it allows. + if prefix not in ("constituency", "local_authority"): return None from policyengine.core.region import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) @@ -509,7 +514,11 @@ def release(reference: ArtifactPathReference) -> tuple[str, str]: msg = f"{matrix_name} entry {year} is not a weight matrix" raise TypeError(msg) matrix_households = weights.shape[-1] - households = len(pd.DataFrame(dataset.data.entity_data["household"])) + entity_data = getattr(dataset.data, "entity_data", None) + if not isinstance(entity_data, dict): + msg = f"the {dataset_selection.name!r} dataset has no entity data loaded" + raise TypeError(msg) + households = len(pd.DataFrame(entity_data["household"])) if households != matrix_households: msg = ( f"UK region {region!r} reweights households with {matrix_name}, " @@ -811,7 +820,7 @@ def _build_simulation( spm_kwargs: dict[str, Any] = {"spm": selection} if selection is not None else {} country = params.get("country", "us") country_module = _country_module(country) - policy = _compile_interval_policy( + runtime_policy = _compile_interval_policy( policy, country_module=country_module, year=_parse_year(params) ) scoping_strategy = region.scoping_strategy @@ -820,7 +829,7 @@ def _build_simulation( params, country=country, dataset_is_default=dataset_selection.is_default, - policy=policy, + policy=runtime_policy, region_code=region.code, scoping_strategy=scoping_strategy, year=_parse_year(params), @@ -835,14 +844,14 @@ def _build_simulation( id=simulation_id, dataset=dataset, tax_benefit_model_version=country_module.model, - policy=policy, + policy=runtime_policy, scoping_strategy=scoping_strategy, ) return StageReportingSimulation( **spm_kwargs, dataset=dataset, tax_benefit_model_version=country_module.model, - policy=policy, + policy=runtime_policy, scoping_strategy=scoping_strategy, ) diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt index 77651b94b..ec72836f6 100644 --- a/rules/baselines/pyright-executor.txt +++ b/rules/baselines/pyright-executor.txt @@ -64,11 +64,6 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/s projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportAttributeAccessIssue :: Cannot access attribute "household" for class "BaseModel" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "household" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "region_type" of type "RegionType" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAssignmentType :: Type "dict[str, Any] | Policy | None" is not assignable to declared type "dict[str, Any] | None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportOptionalMemberAccess :: "items" is not a known attribute of "None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "actual_values" of type "Series" in function "_compare_column" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "expected_values" of type "Series" in function "_compare_column" diff --git a/rules/registry.toml b/rules/registry.toml index e2ac39930..0c792addf 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1143,7 +1143,7 @@ job = "pyright" invocation = "python -m scripts.rules pyright" [rule.baseline] file = "rules/baselines/pyright-executor.txt" -count = 77 +count = 72 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad" good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"