From 99f40f31c22d9b255d175ec0d5a6df6ab854ebd4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:05:11 +0400 Subject: [PATCH 01/14] Publish each model version's entities and basic inputs A model version now records the country package's entities (key, plural, and whether it is the person entity) and its basic inputs, the household fields shown by default. Entity labels, descriptions and roles are not recorded. A country package without basic inputs publishes an empty list. Co-Authored-By: Claude Opus 5.5 --- src/policyengine/core/__init__.py | 1 + src/policyengine/core/entity.py | 17 +++ .../core/tax_benefit_model_version.py | 12 ++ .../common/model_version.py | 14 +++ tests/test_model_entities.py | 112 ++++++++++++++++++ 5 files changed, 156 insertions(+) create mode 100644 src/policyengine/core/entity.py create mode 100644 tests/test_model_entities.py diff --git a/src/policyengine/core/__init__.py b/src/policyengine/core/__init__.py index 70391746..e519b538 100644 --- a/src/policyengine/core/__init__.py +++ b/src/policyengine/core/__init__.py @@ -10,6 +10,7 @@ from .dataset import YearData as YearData from .dataset import map_to_entity as map_to_entity from .dynamic import Dynamic as Dynamic +from .entity import Entity as Entity from .output import Output as Output from .output import OutputCollection as OutputCollection from .parameter import Parameter as Parameter diff --git a/src/policyengine/core/entity.py b/src/policyengine/core/entity.py new file mode 100644 index 00000000..15375b61 --- /dev/null +++ b/src/policyengine/core/entity.py @@ -0,0 +1,17 @@ +from pydantic import BaseModel, Field + + +class Entity(BaseModel): + """One entity of a country model, such as a person or a household. + + Records what a household builder needs to place its inputs: the key that + ``Variable.entity`` names, the plural that groups a household's members of + this entity, and whether it is the person entity. Entity labels, + descriptions and roles are not recorded. + """ + + key: str = Field(description="Entity key, as named by Variable.entity") + plural: str = Field( + description="Plural key grouping a household's entities of this kind" + ) + is_person: bool = Field(description="Whether this is the person entity") diff --git a/src/policyengine/core/tax_benefit_model_version.py b/src/policyengine/core/tax_benefit_model_version.py index 977561c1..dc7a2de4 100644 --- a/src/policyengine/core/tax_benefit_model_version.py +++ b/src/policyengine/core/tax_benefit_model_version.py @@ -13,6 +13,7 @@ ) from policyengine.provenance.trace import build_trace_tro_from_release_bundle +from .entity import Entity from .tax_benefit_model import TaxBenefitModel if TYPE_CHECKING: @@ -46,6 +47,17 @@ class TaxBenefitModelVersion(BaseModel): variables: list["Variable"] = Field(default_factory=list) parameters: list["Parameter"] = Field(default_factory=list) parameter_nodes: list["ParameterNode"] = Field(default_factory=list) + entities: list[Entity] = Field( + default_factory=list, + description="The country model's entities, in the country package's order", + ) + basic_inputs: list[str] = Field( + default_factory=list, + description=( + "Variables a household builder shows by default, as the country " + "package declares them" + ), + ) # Region registry for geographic simulations region_registry: "Optional[RegionRegistry]" = Field( diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index 3b72ad2c..be2a131c 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -30,6 +30,7 @@ import pandas as pd from policyengine.core import ( + Entity, Parameter, ParameterNode, StaleOutputError, @@ -166,6 +167,7 @@ def __init__(self, **kwargs: Any) -> None: self.id = f"{self.model.id}@{self.version}" system = self._load_system() + self._populate_entities(system) self._populate_variables(system) self._populate_parameters(system) @@ -189,6 +191,18 @@ def _dataset_class(self): raise NotImplementedError # --- Shared loading helpers ------------------------------------------ + def _populate_entities(self, system) -> None: + """Copy the entities and basic inputs the country package declares.""" + self.entities = [ + Entity( + key=core_entity.key, + plural=core_entity.plural, + is_person=bool(core_entity.is_person), + ) + for core_entity in system.entities + ] + self.basic_inputs = list(getattr(system, "basic_inputs", None) or []) + def _populate_variables(self, system) -> None: from policyengine_core.enums import Enum from policyengine_core.parameters.operations.get_parameter import ( diff --git a/tests/test_model_entities.py b/tests/test_model_entities.py new file mode 100644 index 00000000..c914a149 --- /dev/null +++ b/tests/test_model_entities.py @@ -0,0 +1,112 @@ +"""Tests for the entities and basic inputs a model version publishes.""" + +import json +from types import SimpleNamespace + +import pytest + +from policyengine.core import Entity, TaxBenefitModel, TaxBenefitModelVersion +from policyengine.tax_benefit_models.common.model_version import ( + MicrosimulationModelVersion, +) +from policyengine.tax_benefit_models.uk import uk_latest +from policyengine.tax_benefit_models.us import us_latest + + +def _model_version() -> TaxBenefitModelVersion: + return TaxBenefitModelVersion( + id="test@1", + model=TaxBenefitModel(id="test", name="Test"), + version="1", + ) + + +def _core_entity(key: str, plural: str, is_person: bool) -> SimpleNamespace: + return SimpleNamespace( + key=key, + plural=plural, + is_person=is_person, + label=key.title(), + doc=f"The {key} entity.", + ) + + +def test_populate_entities_records_key_plural_and_person_flag_in_order(): + model = _model_version() + system = SimpleNamespace( + entities=[ + _core_entity("person", "people", True), + _core_entity("benunit", "benunits", False), + _core_entity("household", "households", False), + ], + basic_inputs=["region", "employment_income", "age"], + ) + + MicrosimulationModelVersion._populate_entities(model, system) + + assert model.entities == [ + Entity(key="person", plural="people", is_person=True), + Entity(key="benunit", plural="benunits", is_person=False), + Entity(key="household", plural="households", is_person=False), + ] + assert model.basic_inputs == ["region", "employment_income", "age"] + + +def test_populate_entities_accepts_a_country_model_without_basic_inputs(): + model = _model_version() + system = SimpleNamespace( + entities=[_core_entity("person", "people", True)], + basic_inputs=None, + ) + + MicrosimulationModelVersion._populate_entities(model, system) + + assert model.basic_inputs == [] + + +def test_model_version_defaults_to_no_entities_or_basic_inputs(): + model = _model_version() + + assert model.entities == [] + assert model.basic_inputs == [] + + +def test_entities_and_basic_inputs_round_trip_through_json(): + model = _model_version() + model.entities = [Entity(key="person", plural="people", is_person=True)] + model.basic_inputs = ["age"] + + payload = json.loads(model.model_dump_json(include={"entities", "basic_inputs"})) + + assert payload == { + "entities": [{"key": "person", "plural": "people", "is_person": True}], + "basic_inputs": ["age"], + } + restored_entities = [ + Entity.model_validate(entity) for entity in payload["entities"] + ] + assert restored_entities == model.entities + + +@pytest.mark.parametrize( + ("model", "country_system_module"), + [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], + ids=["uk", "us"], +) +def test_country_models_publish_their_entities_and_basic_inputs( + model, country_system_module +): + country_system = pytest.importorskip(country_system_module).system + + assert [ + (entity.key, entity.plural, entity.is_person) for entity in model.entities + ] == [ + (entity.key, entity.plural, bool(entity.is_person)) + for entity in country_system.entities + ] + assert [entity.key for entity in model.entities if entity.is_person] == ["person"] + assert model.basic_inputs == list(country_system.basic_inputs) + assert model.basic_inputs, "Each country package declares basic inputs" + assert all(name in model.variables_by_name for name in model.basic_inputs) + entity_keys = {entity.key for entity in model.entities} + assert {variable.entity for variable in model.variables} <= entity_keys From 2fdc802900438dd73d95b8266131eb76c513bfc2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:14:52 +0400 Subject: [PATCH 02/14] Publish whether each variable is entered or hidden, and its module Each variable now records whether users enter it rather than it being calculated (core's is_input_variable), whether it is hidden from input lists, and the module that defines it. Variables built from older country models that do not declare these leave them as None. Co-Authored-By: Claude Opus 5.5 --- src/policyengine/core/variable.py | 19 +++ .../common/model_version.py | 11 ++ tests/test_variable_metadata.py | 115 ++++++++++++++++++ 3 files changed, 145 insertions(+) diff --git a/src/policyengine/core/variable.py b/src/policyengine/core/variable.py index 415081a7..bbab5550 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -30,6 +30,25 @@ class Variable(BaseModel): max_value: Optional[Union[int, float]] = None is_period_size_independent: Optional[bool] = None metadata: dict[str, Any] = Field(default_factory=dict) + is_input_variable: Optional[bool] = Field( + default=None, + description=( + "Whether users enter the variable rather than it being calculated: " + "it has no formula and no adds or subtracts. None when the country " + "model does not say." + ), + ) + hidden_input: Optional[bool] = Field( + default=None, + description=( + "Whether the variable is hidden from lists of inputs. None when the " + "country model does not say." + ), + ) + module_name: Optional[str] = Field( + default=None, + description="Module that defines the variable, as the country model names it", + ) @field_validator("reference", mode="before") @classmethod diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index be2a131c..f03a44be 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -94,6 +94,14 @@ def build_runtime_dataset_provenance( return provenance +def _core_input_flag(core_variable: Any) -> Optional[bool]: + """Return whether users enter a core variable, or None if core cannot say.""" + is_input_variable = getattr(core_variable, "is_input_variable", None) + if is_input_variable is None: + return None + return bool(is_input_variable()) + + class MicrosimulationModelVersion(TaxBenefitModelVersion): """Shared init / save / load logic for country microsim model versions. @@ -239,6 +247,9 @@ def _populate_variables(self, system) -> None: var_obj, "is_period_size_independent", None ), metadata=getattr(var_obj, "metadata", None) or {}, + is_input_variable=_core_input_flag(var_obj), + hidden_input=getattr(var_obj, "hidden_input", None), + module_name=getattr(var_obj, "module_name", None), ) if ( hasattr(var_obj, "possible_values") diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index 4bfa3d5a..4a4e69e2 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -4,6 +4,8 @@ import json from types import SimpleNamespace +import pytest + from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion from policyengine.tax_benefit_models.common.model_version import ( MicrosimulationModelVersion, @@ -142,3 +144,116 @@ def test_us_model_exposes_selected_variable_metadata(): ] assert employment_income.is_period_size_independent is False assert employment_income.metadata == {} + + +def test_populate_variables_records_input_flags_and_module(): + model = _model_version() + entered_input = SimpleNamespace( + name="hours_worked", + label="Hours worked", + entity=SimpleNamespace(key="person"), + documentation=None, + value_type=float, + default_value=0.0, + possible_values=None, + adds=None, + subtracts=None, + is_input_variable=lambda: True, + hidden_input=True, + module_name="household.labor.hours_worked", + ) + calculated_value = SimpleNamespace( + name="total_income", + label="Total income", + entity=SimpleNamespace(key="person"), + documentation=None, + value_type=float, + default_value=0.0, + possible_values=None, + adds=["hours_worked"], + subtracts=None, + is_input_variable=lambda: False, + hidden_input=False, + module_name="household.income.total_income", + ) + system = SimpleNamespace( + variables={ + entered_input.name: entered_input, + calculated_value.name: calculated_value, + }, + parameters=None, + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + hours_worked = model.get_variable("hours_worked") + assert hours_worked.is_input_variable is True + assert hours_worked.hidden_input is True + assert hours_worked.module_name == "household.labor.hours_worked" + total_income = model.get_variable("total_income") + assert total_income.is_input_variable is False + assert total_income.hidden_input is False + assert total_income.module_name == "household.income.total_income" + payload = json.loads( + total_income.model_dump_json( + include={"is_input_variable", "hidden_input", "module_name"} + ) + ) + assert payload == { + "is_input_variable": False, + "hidden_input": False, + "module_name": "household.income.total_income", + } + + +def test_input_flags_and_module_default_to_unrecorded_for_older_country_models(): + model = _model_version() + core_variable = SimpleNamespace( + name="minimal_input", + label="Minimal input", + entity=SimpleNamespace(key="person"), + documentation=None, + value_type=int, + default_value=0, + possible_values=None, + adds=None, + subtracts=None, + ) + system = SimpleNamespace( + variables={core_variable.name: core_variable}, + parameters=None, + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + variable = model.get_variable("minimal_input") + assert variable.is_input_variable is None + assert variable.hidden_input is None + assert variable.module_name is None + + +@pytest.mark.parametrize( + ("model", "country_system_module"), + [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], + ids=["uk", "us"], +) +def test_country_models_publish_input_flags_and_modules_as_core_declares_them( + model, country_system_module +): + country_system = pytest.importorskip(country_system_module).system + + for name, core_variable in country_system.variables.items(): + variable = model.get_variable(name) + assert variable.is_input_variable is core_variable.is_input_variable(), name + assert variable.hidden_input is core_variable.hidden_input, name + assert variable.module_name == core_variable.module_name, name + assert model.get_variable("age").is_input_variable is True + assert any(variable.is_input_variable is False for variable in model.variables) + + +def test_us_model_publishes_hidden_inputs(): + hidden_inputs = [ + variable.name for variable in us_latest.variables if variable.hidden_input + ] + + assert hidden_inputs, "The US package hides some inputs from input lists" From 5c940058a9855f20b8fd4bcda2b9c0a032aa133b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:20:58 +0400 Subject: [PATCH 03/14] Publish the labels of each enum variable's values Enum variables now carry enum_values, each value's stored name with the label users see, in the order the country model declares them. The existing possible_values list of names is unchanged. A label that is not a string, such as an integer, is published as its string form. Co-Authored-By: Claude Opus 5.5 --- src/policyengine/core/__init__.py | 1 + src/policyengine/core/variable.py | 17 +++ .../common/model_version.py | 5 + tests/test_variable_metadata.py | 117 +++++++++++++++++- 4 files changed, 139 insertions(+), 1 deletion(-) diff --git a/src/policyengine/core/__init__.py b/src/policyengine/core/__init__.py index e519b538..b30b34fb 100644 --- a/src/policyengine/core/__init__.py +++ b/src/policyengine/core/__init__.py @@ -32,6 +32,7 @@ from .tax_benefit_model_version import ( TaxBenefitModelVersion as TaxBenefitModelVersion, ) +from .variable import EnumValue as EnumValue from .variable import Variable as Variable # Rebuild models to resolve forward references diff --git a/src/policyengine/core/variable.py b/src/policyengine/core/variable.py index bbab5550..5412490c 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -8,6 +8,15 @@ from .tax_benefit_model_version import TaxBenefitModelVersion +class EnumValue(BaseModel): + """One value an enum variable can take, with the label users see.""" + + value: str = Field( + description="Name stored in household data, such as 'SINGLE'", + ) + label: str = Field(description="Label shown to users, such as 'Single'") + + class Variable(BaseModel): id: str name: str @@ -49,6 +58,14 @@ class Variable(BaseModel): default=None, description="Module that defines the variable, as the country model names it", ) + enum_values: Optional[list[EnumValue]] = Field( + default=None, + description=( + "For an enum variable, each value it can take with its label, in " + "the order the country model declares them. possible_values holds " + "the same names without labels." + ), + ) @field_validator("reference", mode="before") @classmethod diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index f03a44be..53404cc7 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -31,6 +31,7 @@ from policyengine.core import ( Entity, + EnumValue, Parameter, ParameterNode, StaleOutputError, @@ -261,6 +262,10 @@ def _populate_variables(self, system) -> None: var_obj.possible_values._value2member_map_.values(), ) ) + variable.enum_values = [ + EnumValue(value=member.name, label=str(member.value)) + for member in var_obj.possible_values + ] # Resolve parameter-path adds/subtracts to concrete lists so # consumers always see list[str]. for attr in ("adds", "subtracts"): diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index 4a4e69e2..837b6a4e 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -2,11 +2,12 @@ import datetime import json +from enum import Enum from types import SimpleNamespace import pytest -from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion +from policyengine.core import EnumValue, TaxBenefitModel, TaxBenefitModelVersion from policyengine.tax_benefit_models.common.model_version import ( MicrosimulationModelVersion, ) @@ -257,3 +258,117 @@ def test_us_model_publishes_hidden_inputs(): ] assert hidden_inputs, "The US package hides some inputs from input lists" + + +class _FilingStatus(Enum): + SINGLE = "Single" + JOINT = "Joint" + HEAD_OF_HOUSEHOLD = "Head of household" + + +class _NumberedStatus(Enum): + FIRST = 1 + SECOND = 2 + + +def _enum_core_variable(name: str, possible_values: type[Enum]) -> SimpleNamespace: + first_member = next(iter(possible_values)) + return SimpleNamespace( + name=name, + label=name.replace("_", " ").capitalize(), + entity=SimpleNamespace(key="tax_unit"), + documentation=None, + value_type=str, + default_value=first_member.name, + possible_values=possible_values, + adds=None, + subtracts=None, + ) + + +def test_populate_variables_records_enum_values_with_their_labels_in_order(): + model = _model_version() + filing_status = _enum_core_variable("filing_status", _FilingStatus) + numbered_status = _enum_core_variable("numbered_status", _NumberedStatus) + system = SimpleNamespace( + variables={ + filing_status.name: filing_status, + numbered_status.name: numbered_status, + }, + parameters=None, + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + variable = model.get_variable("filing_status") + assert variable.possible_values == ["SINGLE", "JOINT", "HEAD_OF_HOUSEHOLD"] + assert variable.enum_values == [ + EnumValue(value="SINGLE", label="Single"), + EnumValue(value="JOINT", label="Joint"), + EnumValue(value="HEAD_OF_HOUSEHOLD", label="Head of household"), + ] + assert json.loads(variable.model_dump_json(include={"enum_values"})) == { + "enum_values": [ + {"value": "SINGLE", "label": "Single"}, + {"value": "JOINT", "label": "Joint"}, + {"value": "HEAD_OF_HOUSEHOLD", "label": "Head of household"}, + ] + } + assert [ + (enum_value.value, enum_value.label) + for enum_value in model.get_variable("numbered_status").enum_values + ] == [("FIRST", "1"), ("SECOND", "2")] + + +def test_variables_without_possible_values_have_no_enum_values(): + model = _model_version() + core_variable = SimpleNamespace( + name="minimal_input", + label="Minimal input", + entity=SimpleNamespace(key="person"), + documentation=None, + value_type=int, + default_value=0, + adds=None, + subtracts=None, + ) + system = SimpleNamespace( + variables={core_variable.name: core_variable}, + parameters=None, + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + assert model.get_variable("minimal_input").enum_values is None + + +@pytest.mark.parametrize( + ("model", "country_system_module"), + [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], + ids=["uk", "us"], +) +def test_country_models_publish_every_enum_label(model, country_system_module): + country_system = pytest.importorskip(country_system_module).system + + enum_variable_count = 0 + for name, core_variable in country_system.variables.items(): + possible_values = getattr(core_variable, "possible_values", None) + variable = model.get_variable(name) + if possible_values is None: + assert variable.enum_values is None, name + continue + enum_variable_count += 1 + assert [ + (enum_value.value, enum_value.label) for enum_value in variable.enum_values + ] == [(member.name, str(member.value)) for member in possible_values], name + assert [enum_value.value for enum_value in variable.enum_values] == ( + variable.possible_values + ), name + assert enum_variable_count > 0 + + +def test_us_filing_status_publishes_its_labels(): + filing_status = us_latest.get_variable("filing_status") + + assert EnumValue(value="SINGLE", label="Single") in filing_status.enum_values + assert EnumValue(value="JOINT", label="Joint") in filing_status.enum_values From 0d36dd161b2e3662439ba9180a6a4c5244c84d90 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:24:26 +0400 Subject: [PATCH 04/14] Publish the parameter path that adds or subtracts names Where a variable's adds or subtracts names a parameter instead of a list, the variable now records that path in adds_parameter_path or subtracts_parameter_path, so a consumer can read the list for any year. adds and subtracts keep the list resolved as of 2025-01-01. Co-Authored-By: Claude Opus 5.5 --- src/policyengine/core/variable.py | 17 +++ .../common/model_version.py | 4 +- tests/test_variable_metadata.py | 131 ++++++++++++++++++ 3 files changed, 151 insertions(+), 1 deletion(-) diff --git a/src/policyengine/core/variable.py b/src/policyengine/core/variable.py index 5412490c..d4096959 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -58,6 +58,23 @@ class Variable(BaseModel): default=None, description="Module that defines the variable, as the country model names it", ) + adds_parameter_path: Optional[str] = Field( + default=None, + description=( + "Path of the parameter whose value lists the variables this one " + "adds, when the country model names a parameter rather than a " + "list. adds then holds that parameter's list as of 2025-01-01 only." + ), + ) + subtracts_parameter_path: Optional[str] = Field( + default=None, + description=( + "Path of the parameter whose value lists the variables this one " + "subtracts, when the country model names a parameter rather than a " + "list. subtracts then holds that parameter's list as of 2025-01-01 " + "only." + ), + ) enum_values: Optional[list[EnumValue]] = Field( default=None, description=( diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index 53404cc7..fb88873e 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -267,12 +267,14 @@ def _populate_variables(self, system) -> None: for member in var_obj.possible_values ] # Resolve parameter-path adds/subtracts to concrete lists so - # consumers always see list[str]. + # consumers always see list[str], and record the path itself, + # since the parameter's list can differ in other years. for attr in ("adds", "subtracts"): value = getattr(var_obj, attr, None) if value is None: continue if isinstance(value, str): + setattr(variable, f"{attr}_parameter_path", value) try: param = get_parameter(system.parameters, value) setattr(variable, attr, list(param("2025-01-01"))) diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index 837b6a4e..f0d563c2 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import pytest +from policyengine_core.parameters import ParameterNode as CoreParameterNode from policyengine.core import EnumValue, TaxBenefitModel, TaxBenefitModelVersion from policyengine.tax_benefit_models.common.model_version import ( @@ -372,3 +373,133 @@ def test_us_filing_status_publishes_its_labels(): assert EnumValue(value="SINGLE", label="Single") in filing_status.enum_values assert EnumValue(value="JOINT", label="Joint") in filing_status.enum_values + + +def _components_parameter_tree() -> CoreParameterNode: + return CoreParameterNode( + "", + data={ + "gov": { + "example": { + "components": { + "description": "Income components, which change in 2026", + "values": { + "2020-01-01": ["wages"], + "2026-01-01": ["wages", "tips"], + }, + }, + "deductions": { + "description": "Deductions", + "values": {"2020-01-01": ["pension_contributions"]}, + }, + } + } + }, + ) + + +def _summed_core_variable(name: str, adds, subtracts) -> SimpleNamespace: + return SimpleNamespace( + name=name, + label=name.replace("_", " ").capitalize(), + entity=SimpleNamespace(key="person"), + documentation=None, + value_type=float, + default_value=0.0, + possible_values=None, + adds=adds, + subtracts=subtracts, + ) + + +def test_populate_variables_records_parameter_paths_named_by_adds_and_subtracts(): + model = _model_version() + parameter_named = _summed_core_variable( + "net_income", "gov.example.components", "gov.example.deductions" + ) + variable_named = _summed_core_variable("gross_income", ["wages", "tips"], None) + system = SimpleNamespace( + variables={ + parameter_named.name: parameter_named, + variable_named.name: variable_named, + }, + parameters=_components_parameter_tree(), + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + net_income = model.get_variable("net_income") + assert net_income.adds_parameter_path == "gov.example.components" + assert net_income.subtracts_parameter_path == "gov.example.deductions" + # The lists resolved as of 2025-01-01 are unchanged for existing users. + assert net_income.adds == ["wages"] + assert net_income.subtracts == ["pension_contributions"] + assert json.loads( + net_income.model_dump_json( + include={"adds_parameter_path", "subtracts_parameter_path"} + ) + ) == { + "adds_parameter_path": "gov.example.components", + "subtracts_parameter_path": "gov.example.deductions", + } + gross_income = model.get_variable("gross_income") + assert gross_income.adds == ["wages", "tips"] + assert gross_income.adds_parameter_path is None + assert gross_income.subtracts_parameter_path is None + + +def test_parameter_path_is_published_when_the_parameter_cannot_be_resolved(): + model = _model_version() + core_variable = _summed_core_variable("net_income", "gov.example.missing", None) + system = SimpleNamespace( + variables={core_variable.name: core_variable}, + parameters=_components_parameter_tree(), + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + variable = model.get_variable("net_income") + assert variable.adds is None + assert variable.adds_parameter_path == "gov.example.missing" + + +@pytest.mark.parametrize( + ("model", "country_system_module"), + [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], + ids=["uk", "us"], +) +def test_country_models_publish_every_parameter_named_by_adds_or_subtracts( + model, country_system_module +): + country_system = pytest.importorskip(country_system_module).system + + for name, core_variable in country_system.variables.items(): + variable = model.get_variable(name) + for attribute in ("adds", "subtracts"): + declared = getattr(core_variable, attribute, None) + published_path = getattr(variable, f"{attribute}_parameter_path") + if isinstance(declared, str): + assert published_path == declared, (name, attribute) + assert published_path in model.parameters_by_name, (name, attribute) + else: + assert published_path is None, (name, attribute) + + +def test_us_model_publishes_parameter_named_parts(): + adds_paths = [ + variable for variable in us_latest.variables if variable.adds_parameter_path + ] + subtracts_paths = [ + variable + for variable in us_latest.variables + if variable.subtracts_parameter_path + ] + + assert len(adds_paths) >= 200 + assert len(subtracts_paths) >= 1 + household_state_benefits = us_latest.get_variable("household_state_benefits") + assert ( + household_state_benefits.adds_parameter_path + == "gov.household.household_state_benefits" + ) + assert isinstance(household_state_benefits.adds, list) From dc49f18c3ebe9e334f89c22a818143b4a60d54a5 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:27:57 +0400 Subject: [PATCH 05/14] Publish each parameter's economy and household flags Each parameter now carries the economy and household flags from its country package metadata, True when the metadata omits them, as the Public API's v1 metadata does. Parameters built without a country package default to True. Co-Authored-By: Claude Opus 5.5 --- src/policyengine/core/parameter.py | 16 +++ .../common/model_version.py | 2 + tests/test_parameter_metadata.py | 97 +++++++++++++++++++ 3 files changed, 115 insertions(+) create mode 100644 tests/test_parameter_metadata.py diff --git a/src/policyengine/core/parameter.py b/src/policyengine/core/parameter.py index 680120cf..5f61ba51 100644 --- a/src/policyengine/core/parameter.py +++ b/src/policyengine/core/parameter.py @@ -21,6 +21,22 @@ class Parameter(BaseModel): data_type: Optional[type] = None tax_benefit_model_version: TaxBenefitModelVersion unit: Optional[str] = None + economy: bool = Field( + default=True, + description=( + "Whether the parameter applies to economy-wide (population) " + "analysis, from the country package's 'economy' metadata. True " + "when the package does not say." + ), + ) + household: bool = Field( + default=True, + description=( + "Whether the parameter applies to household calculations, from " + "the country package's 'household' metadata. True when the " + "package does not say." + ), + ) # Lazy loading: store core param ref, build values on demand _core_param: Any = PrivateAttr(default=None) diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index fb88873e..b1adf88d 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -302,6 +302,8 @@ def _populate_parameters(self, system) -> None: description=param_node.description, data_type=type(param_node(2025)), unit=param_node.metadata.get("unit"), + economy=bool(param_node.metadata.get("economy", True)), + household=bool(param_node.metadata.get("household", True)), _core_param=param_node, ) self.add_parameter(parameter) diff --git a/tests/test_parameter_metadata.py b/tests/test_parameter_metadata.py new file mode 100644 index 00000000..a273a3ea --- /dev/null +++ b/tests/test_parameter_metadata.py @@ -0,0 +1,97 @@ +"""Tests for the economy and household flags each parameter publishes.""" + +import json +from types import SimpleNamespace + +import pytest +from policyengine_core.parameters import Parameter as CoreParameter +from policyengine_core.parameters import ParameterNode as CoreParameterNode + +from policyengine.core import Parameter, TaxBenefitModel, TaxBenefitModelVersion +from policyengine.tax_benefit_models.common.model_version import ( + MicrosimulationModelVersion, +) +from policyengine.tax_benefit_models.uk import uk_latest +from policyengine.tax_benefit_models.us import us_latest + + +def _model_version() -> TaxBenefitModelVersion: + return TaxBenefitModelVersion( + id="test@1", + model=TaxBenefitModel(id="test", name="Test"), + version="1", + ) + + +def _flagged_parameter(metadata: dict) -> dict: + return { + "description": "An example amount", + "values": {"2020-01-01": 100}, + "metadata": metadata, + } + + +def test_populate_parameters_records_economy_and_household_flags(): + model = _model_version() + parameters = CoreParameterNode( + "", + data={ + "gov": { + "example": { + "economy_only": _flagged_parameter({"household": False}), + "household_only": _flagged_parameter({"economy": False}), + "unflagged": _flagged_parameter({}), + } + } + }, + ) + system = SimpleNamespace(parameters=parameters, variables={}) + + MicrosimulationModelVersion._populate_parameters(model, system) + + economy_only = model.get_parameter("gov.example.economy_only") + assert (economy_only.economy, economy_only.household) == (True, False) + household_only = model.get_parameter("gov.example.household_only") + assert (household_only.economy, household_only.household) == (False, True) + # As in the country packages' metadata, an absent flag means the + # parameter is offered. + unflagged = model.get_parameter("gov.example.unflagged") + assert (unflagged.economy, unflagged.household) == (True, True) + assert json.loads( + household_only.model_dump_json(include={"economy", "household"}) + ) == { + "economy": False, + "household": True, + } + + +def test_parameter_flags_default_to_offered(): + model = _model_version() + manually_built = Parameter( + name="gov.example.amount", tax_benefit_model_version=model + ) + + assert manually_built.economy is True + assert manually_built.household is True + + +@pytest.mark.parametrize( + ("model", "country_system_module"), + [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], + ids=["uk", "us"], +) +def test_country_models_publish_each_parameters_flags(model, country_system_module): + country_system = pytest.importorskip(country_system_module).system + + for core_parameter in country_system.parameters.get_descendants(): + if not isinstance(core_parameter, CoreParameter): + continue + parameter = model.get_parameter(core_parameter.name) + assert parameter.economy is bool( + core_parameter.metadata.get("economy", True) + ), core_parameter.name + assert parameter.household is bool( + core_parameter.metadata.get("household", True) + ), core_parameter.name + assert any(not parameter.economy for parameter in model.parameters) + assert any(not parameter.household for parameter in model.parameters) From 89334ee98ffee5c7a1988d9af81ac373fda4480b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:54:10 +0400 Subject: [PATCH 06/14] Publish each bracket schedule as a parameter node with its label A bracket schedule, such as gov.hmrc.income_tax.rates.uk, is now a parameter node carrying its stored label (None when the country package has none) and description, listed before its brackets. Its brackets (name[0], name[1], ...) stay nodes of their own, and their rate, threshold and amount parameters are unchanged. Schedules with no brackets are left out, as empty nodes already are. Co-Authored-By: Claude Opus 5.5 --- src/policyengine/core/parameter_node.py | 4 + .../common/model_version.py | 14 ++++ tests/test_parameter_nodes.py | 82 +++++++++++++++++++ 3 files changed, 100 insertions(+) diff --git a/src/policyengine/core/parameter_node.py b/src/policyengine/core/parameter_node.py index 54d384a5..af060b79 100644 --- a/src/policyengine/core/parameter_node.py +++ b/src/policyengine/core/parameter_node.py @@ -14,6 +14,10 @@ class ParameterNode(BaseModel): "gov.hmrc", "gov.hmrc.income_tax"). They provide structure and human-readable labels for navigating the parameter tree, but don't have values themselves. + A bracket schedule (e.g., "gov.hmrc.income_tax.rates.uk") is a node, and + so is each of its brackets ("gov.hmrc.income_tax.rates.uk[0]"), whose + parameters are its rate, threshold or amount. + Unlike Parameter objects (which are leaf nodes with actual values), ParameterNode objects are purely organizational. """ diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index b1adf88d..1673d59f 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -287,6 +287,7 @@ def _populate_variables(self, system) -> None: def _populate_parameters(self, system) -> None: from policyengine_core.parameters import Parameter as CoreParameter from policyengine_core.parameters import ParameterNode as CoreParameterNode + from policyengine_core.parameters import ParameterScale as CoreParameterScale scale_lookup = build_scale_lookup(system) @@ -307,6 +308,19 @@ def _populate_parameters(self, system) -> None: _core_param=param_node, ) self.add_parameter(parameter) + elif isinstance(param_node, CoreParameterScale): + # A bracket schedule is not a core ParameterNode, but it groups + # its brackets (``name[0]``, ``name[1]``, ...) as a node does. + if not param_node.brackets: + continue + schedule = ParameterNode( + id=self.id + "-" + param_node.name, + name=param_node.name, + label=param_node.metadata.get("label"), + description=param_node.description, + tax_benefit_model_version=self, + ) + self.add_parameter_node(schedule) elif isinstance(param_node, CoreParameterNode): if not param_node.children: continue diff --git a/tests/test_parameter_nodes.py b/tests/test_parameter_nodes.py index 99a8edd6..eced690f 100644 --- a/tests/test_parameter_nodes.py +++ b/tests/test_parameter_nodes.py @@ -4,6 +4,7 @@ import pytest from policyengine_core.parameters import ParameterNode as CoreParameterNode +from policyengine_core.parameters import ParameterScale as CoreParameterScale from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion from policyengine.tax_benefit_models.common.model_version import ( @@ -54,3 +55,84 @@ def test_country_models_do_not_expose_python_cache_nodes(model): assert all( "__pycache__" not in name.split(".") for name in model.parameter_nodes_by_name ) + + +def _bracket(rate: float, threshold: float) -> dict: + return { + "rate": {"values": {"2020-01-01": rate}}, + "threshold": {"values": {"2020-01-01": threshold}}, + } + + +def _system_with_bracket_schedules() -> SimpleNamespace: + parameters = CoreParameterNode( + "", + data={ + "gov": { + "tax": { + "rates": { + "description": "Income tax rates by band", + "metadata": { + "label": "Income tax rates", + "type": "marginal_rate", + }, + "brackets": [_bracket(0.1, 0), _bracket(0.2, 1_000)], + }, + "unlabelled_rates": { + "description": "A schedule with no stored label", + "metadata": {"type": "marginal_rate"}, + "brackets": [_bracket(0.05, 0)], + }, + } + } + }, + ) + return SimpleNamespace(parameters=parameters, variables={}) + + +def test_populate_parameters_publishes_bracket_schedules_as_labelled_nodes(): + model = _model_version() + + MicrosimulationModelVersion._populate_parameters( + model, _system_with_bracket_schedules() + ) + + schedule = model.get_parameter_node("gov.tax.rates") + assert schedule.label == "Income tax rates" + assert schedule.description == "Income tax rates by band" + assert schedule.id == "test@1-gov.tax.rates" + assert model.get_parameter_node("gov.tax.unlabelled_rates").label is None + node_names = [node.name for node in model.parameter_nodes] + # Each schedule precedes its brackets, which stay nodes of their own. + assert node_names.index("gov.tax.rates") < node_names.index("gov.tax.rates[0]") + assert {"gov.tax.rates[0]", "gov.tax.rates[1]"} <= set(node_names) + assert model.get_parameter("gov.tax.rates[1].rate").label == ( + "Income tax rates (bracket 2 rate)" + ) + + +@pytest.mark.parametrize( + ("model", "country_system_module"), + [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], + ids=["uk", "us"], +) +def test_country_models_publish_every_bracket_schedule(model, country_system_module): + country_system = pytest.importorskip(country_system_module).system + + schedules = [ + core_parameter + for core_parameter in country_system.parameters.get_descendants() + if isinstance(core_parameter, CoreParameterScale) and core_parameter.brackets + ] + assert schedules + for core_schedule in schedules: + schedule = model.get_parameter_node(core_schedule.name) + assert schedule.label == core_schedule.metadata.get("label"), core_schedule.name + assert schedule.description == core_schedule.description + + +def test_uk_income_tax_schedule_is_a_labelled_node(): + schedule = uk_latest.get_parameter_node("gov.hmrc.income_tax.rates.uk") + + assert schedule.label + assert "gov.hmrc.income_tax.rates.uk[0]" in uk_latest.parameter_nodes_by_name From 185b9a3074c255def57964830c67587f0379ee32 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:07:18 +0400 Subject: [PATCH 07/14] Publish each enum value's declared value as its label, unchanged An enum member declared with an integer, such as the US ILPersonalExemptionEligibilityStatus values 1, 2 and 3, now publishes the integer as its label instead of its string form, as the Public API's /{country}/metadata does. EnumValue.label holds a string or an integer. Co-Authored-By: Claude Opus 5.5 --- src/policyengine/core/variable.py | 8 ++- .../common/model_version.py | 2 +- tests/test_variable_metadata.py | 52 +++++++++++++++---- 3 files changed, 50 insertions(+), 12 deletions(-) diff --git a/src/policyengine/core/variable.py b/src/policyengine/core/variable.py index d4096959..79a4bcef 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -14,7 +14,13 @@ class EnumValue(BaseModel): value: str = Field( description="Name stored in household data, such as 'SINGLE'", ) - label: str = Field(description="Label shown to users, such as 'Single'") + label: Union[str, int] = Field( + description=( + "Label shown to users, such as 'Single': the value the country " + "model declares for the member, unchanged, so a member declared " + "with an integer has an integer label" + ), + ) class Variable(BaseModel): diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index 1673d59f..21c27b8d 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -263,7 +263,7 @@ def _populate_variables(self, system) -> None: ) ) variable.enum_values = [ - EnumValue(value=member.name, label=str(member.value)) + EnumValue(value=member.name, label=member.value) for member in var_obj.possible_values ] # Resolve parameter-path adds/subtracts to concrete lists so diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index f0d563c2..dd4160c1 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -290,12 +290,8 @@ def _enum_core_variable(name: str, possible_values: type[Enum]) -> SimpleNamespa def test_populate_variables_records_enum_values_with_their_labels_in_order(): model = _model_version() filing_status = _enum_core_variable("filing_status", _FilingStatus) - numbered_status = _enum_core_variable("numbered_status", _NumberedStatus) system = SimpleNamespace( - variables={ - filing_status.name: filing_status, - numbered_status.name: numbered_status, - }, + variables={filing_status.name: filing_status}, parameters=None, ) @@ -315,10 +311,34 @@ def test_populate_variables_records_enum_values_with_their_labels_in_order(): {"value": "HEAD_OF_HOUSEHOLD", "label": "Head of household"}, ] } - assert [ - (enum_value.value, enum_value.label) - for enum_value in model.get_variable("numbered_status").enum_values - ] == [("FIRST", "1"), ("SECOND", "2")] + + +def test_enum_values_keep_integer_labels_as_declared(): + # The Public API's /{country}/metadata publishes an enum member's + # declared value unchanged, so an integer stays an integer. + model = _model_version() + numbered_status = _enum_core_variable("numbered_status", _NumberedStatus) + system = SimpleNamespace( + variables={numbered_status.name: numbered_status}, + parameters=None, + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + enum_values = model.get_variable("numbered_status").enum_values + assert [(enum_value.value, enum_value.label) for enum_value in enum_values] == [ + ("FIRST", 1), + ("SECOND", 2), + ] + assert all(type(enum_value.label) is int for enum_value in enum_values) + assert json.loads( + model.get_variable("numbered_status").model_dump_json(include={"enum_values"}) + ) == { + "enum_values": [ + {"value": "FIRST", "label": 1}, + {"value": "SECOND", "label": 2}, + ] + } def test_variables_without_possible_values_have_no_enum_values(): @@ -361,7 +381,10 @@ def test_country_models_publish_every_enum_label(model, country_system_module): enum_variable_count += 1 assert [ (enum_value.value, enum_value.label) for enum_value in variable.enum_values - ] == [(member.name, str(member.value)) for member in possible_values], name + ] == [(member.name, member.value) for member in possible_values], name + assert [type(enum_value.label) for enum_value in variable.enum_values] == [ + type(member.value) for member in possible_values + ], name assert [enum_value.value for enum_value in variable.enum_values] == ( variable.possible_values ), name @@ -375,6 +398,15 @@ def test_us_filing_status_publishes_its_labels(): assert EnumValue(value="JOINT", label="Joint") in filing_status.enum_values +def test_us_integer_valued_enum_publishes_integer_labels(): + eligibility = us_latest.get_variable("il_personal_exemption_eligibility_status") + + assert [ + (enum_value.value, enum_value.label) for enum_value in eligibility.enum_values + ] == [("BOTH_ELIGIBLE", 1), ("PARTIALLY_ELIGIBLE", 2), ("NOT_ELIGIBLE", 3)] + assert all(type(enum_value.label) is int for enum_value in eligibility.enum_values) + + def _components_parameter_tree() -> CoreParameterNode: return CoreParameterNode( "", From ad2058862597292abb8021640440e41bb15959c1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:11:13 +0400 Subject: [PATCH 08/14] Skip the new country-model tests when the country package is missing The entity and parameter-flag modules imported uk_latest and us_latest at module level, which the country modules define only when their package is installed, so a missing package failed collection, stub-only tests included, and the importorskip calls inside the tests could never skip. A fixture now skips first and imports the model afterwards. In the variable and parameter-node modules, whose module-level imports predate this change, the added tests import the core system directly instead of appearing to skip. Co-Authored-By: Claude Opus 5.5 --- tests/test_model_entities.py | 35 ++++++++++++++++++++++---------- tests/test_parameter_metadata.py | 33 ++++++++++++++++++++++-------- tests/test_parameter_nodes.py | 3 ++- tests/test_variable_metadata.py | 7 ++++--- 4 files changed, 54 insertions(+), 24 deletions(-) diff --git a/tests/test_model_entities.py b/tests/test_model_entities.py index c914a149..f821ebca 100644 --- a/tests/test_model_entities.py +++ b/tests/test_model_entities.py @@ -1,5 +1,6 @@ """Tests for the entities and basic inputs a model version publishes.""" +import importlib import json from types import SimpleNamespace @@ -9,8 +10,6 @@ from policyengine.tax_benefit_models.common.model_version import ( MicrosimulationModelVersion, ) -from policyengine.tax_benefit_models.uk import uk_latest -from policyengine.tax_benefit_models.us import us_latest def _model_version() -> TaxBenefitModelVersion: @@ -21,6 +20,26 @@ def _model_version() -> TaxBenefitModelVersion: ) +@pytest.fixture(params=["uk", "us"]) +def country(request) -> SimpleNamespace: + """The country's published model version and its core tax-benefit system. + + Skips when the country package is not installed. The model is imported + here, after that check, because the country module defines it only when + the package is installed. + """ + country_id = request.param + country_system = pytest.importorskip(f"policyengine_{country_id}.system").system + country_module = importlib.import_module( + f"policyengine.tax_benefit_models.{country_id}" + ) + return SimpleNamespace( + id=country_id, + model=getattr(country_module, f"{country_id}_latest"), + system=country_system, + ) + + def _core_entity(key: str, plural: str, is_person: bool) -> SimpleNamespace: return SimpleNamespace( key=key, @@ -88,15 +107,9 @@ def test_entities_and_basic_inputs_round_trip_through_json(): assert restored_entities == model.entities -@pytest.mark.parametrize( - ("model", "country_system_module"), - [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], - ids=["uk", "us"], -) -def test_country_models_publish_their_entities_and_basic_inputs( - model, country_system_module -): - country_system = pytest.importorskip(country_system_module).system +def test_country_models_publish_their_entities_and_basic_inputs(country): + model = country.model + country_system = country.system assert [ (entity.key, entity.plural, entity.is_person) for entity in model.entities diff --git a/tests/test_parameter_metadata.py b/tests/test_parameter_metadata.py index a273a3ea..6d07f50a 100644 --- a/tests/test_parameter_metadata.py +++ b/tests/test_parameter_metadata.py @@ -1,5 +1,6 @@ """Tests for the economy and household flags each parameter publishes.""" +import importlib import json from types import SimpleNamespace @@ -11,8 +12,6 @@ from policyengine.tax_benefit_models.common.model_version import ( MicrosimulationModelVersion, ) -from policyengine.tax_benefit_models.uk import uk_latest -from policyengine.tax_benefit_models.us import us_latest def _model_version() -> TaxBenefitModelVersion: @@ -23,6 +22,26 @@ def _model_version() -> TaxBenefitModelVersion: ) +@pytest.fixture(params=["uk", "us"]) +def country(request) -> SimpleNamespace: + """The country's published model version and its core tax-benefit system. + + Skips when the country package is not installed. The model is imported + here, after that check, because the country module defines it only when + the package is installed. + """ + country_id = request.param + country_system = pytest.importorskip(f"policyengine_{country_id}.system").system + country_module = importlib.import_module( + f"policyengine.tax_benefit_models.{country_id}" + ) + return SimpleNamespace( + id=country_id, + model=getattr(country_module, f"{country_id}_latest"), + system=country_system, + ) + + def _flagged_parameter(metadata: dict) -> dict: return { "description": "An example amount", @@ -75,13 +94,9 @@ def test_parameter_flags_default_to_offered(): assert manually_built.household is True -@pytest.mark.parametrize( - ("model", "country_system_module"), - [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], - ids=["uk", "us"], -) -def test_country_models_publish_each_parameters_flags(model, country_system_module): - country_system = pytest.importorskip(country_system_module).system +def test_country_models_publish_each_parameters_flags(country): + model = country.model + country_system = country.system for core_parameter in country_system.parameters.get_descendants(): if not isinstance(core_parameter, CoreParameter): diff --git a/tests/test_parameter_nodes.py b/tests/test_parameter_nodes.py index eced690f..f2c4c3d4 100644 --- a/tests/test_parameter_nodes.py +++ b/tests/test_parameter_nodes.py @@ -1,5 +1,6 @@ """Tests for public parameter-node projection from country models.""" +import importlib from types import SimpleNamespace import pytest @@ -117,7 +118,7 @@ def test_populate_parameters_publishes_bracket_schedules_as_labelled_nodes(): ids=["uk", "us"], ) def test_country_models_publish_every_bracket_schedule(model, country_system_module): - country_system = pytest.importorskip(country_system_module).system + country_system = importlib.import_module(country_system_module).system schedules = [ core_parameter diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index dd4160c1..e9e2d094 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -1,6 +1,7 @@ """Tests for public variable metadata copied from country models.""" import datetime +import importlib import json from enum import Enum from types import SimpleNamespace @@ -242,7 +243,7 @@ def test_input_flags_and_module_default_to_unrecorded_for_older_country_models() def test_country_models_publish_input_flags_and_modules_as_core_declares_them( model, country_system_module ): - country_system = pytest.importorskip(country_system_module).system + country_system = importlib.import_module(country_system_module).system for name, core_variable in country_system.variables.items(): variable = model.get_variable(name) @@ -369,7 +370,7 @@ def test_variables_without_possible_values_have_no_enum_values(): ids=["uk", "us"], ) def test_country_models_publish_every_enum_label(model, country_system_module): - country_system = pytest.importorskip(country_system_module).system + country_system = importlib.import_module(country_system_module).system enum_variable_count = 0 for name, core_variable in country_system.variables.items(): @@ -503,7 +504,7 @@ def test_parameter_path_is_published_when_the_parameter_cannot_be_resolved(): def test_country_models_publish_every_parameter_named_by_adds_or_subtracts( model, country_system_module ): - country_system = pytest.importorskip(country_system_module).system + country_system = importlib.import_module(country_system_module).system for name, core_variable in country_system.variables.items(): variable = model.get_variable(name) From 0fb97b57def8db889b7bfdc613d1b196e49b4a24 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:12:15 +0400 Subject: [PATCH 09/14] Test that bracket schedules without brackets are left out A schedule whose brackets list is empty is not published as a parameter node, as empty organizational nodes already are not. Co-Authored-By: Claude Opus 5.5 --- tests/test_parameter_nodes.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_parameter_nodes.py b/tests/test_parameter_nodes.py index f2c4c3d4..a0060ebc 100644 --- a/tests/test_parameter_nodes.py +++ b/tests/test_parameter_nodes.py @@ -112,6 +112,37 @@ def test_populate_parameters_publishes_bracket_schedules_as_labelled_nodes(): ) +def test_populate_parameters_leaves_out_bracket_schedules_without_brackets(): + parameters = CoreParameterNode( + "", + data={ + "gov": { + "tax": { + "empty_rates": { + "description": "A schedule whose brackets are not set yet", + "metadata": {"label": "Empty rates", "type": "marginal_rate"}, + "brackets": [], + }, + "rates": { + "metadata": {"label": "Income tax rates"}, + "brackets": [_bracket(0.1, 0)], + }, + } + } + }, + ) + assert isinstance(parameters.gov.tax.empty_rates, CoreParameterScale) + model = _model_version() + + MicrosimulationModelVersion._populate_parameters( + model, SimpleNamespace(parameters=parameters, variables={}) + ) + + assert "gov.tax.empty_rates" not in model.parameter_nodes_by_name + assert all(node.name != "gov.tax.empty_rates" for node in model.parameter_nodes) + assert "gov.tax.rates" in model.parameter_nodes_by_name + + @pytest.mark.parametrize( ("model", "country_system_module"), [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], From c28fbb6f214159da30b919815e25139900fbc099 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:20:16 +0400 Subject: [PATCH 10/14] Pin concrete UK and US metadata values in the country-model tests The country tests computed their expected values with the same expressions as the implementation, so they could not detect a difference from what the country packages declare. They now also check fixed values read from the pinned country packages: each country's entity keys, plurals and person flag; economy and household flags of parameters declared false, directly or through propagated metadata, and of one unflagged parameter; input, hidden-input and module values of known variables; the count of parameter-named adds (UK 4, US 237) and subtracts (UK 0, US 2); and stored schedule labels, including schedules without one, which publish None. Co-Authored-By: Claude Opus 5.5 --- tests/test_model_entities.py | 27 +++++++++++++++++++++ tests/test_parameter_metadata.py | 28 ++++++++++++++++++++++ tests/test_parameter_nodes.py | 20 +++++++++++++++- tests/test_variable_metadata.py | 40 ++++++++++++++++++++++++++++++++ 4 files changed, 114 insertions(+), 1 deletion(-) diff --git a/tests/test_model_entities.py b/tests/test_model_entities.py index f821ebca..cba548f2 100644 --- a/tests/test_model_entities.py +++ b/tests/test_model_entities.py @@ -123,3 +123,30 @@ def test_country_models_publish_their_entities_and_basic_inputs(country): assert all(name in model.variables_by_name for name in model.basic_inputs) entity_keys = {entity.key for entity in model.entities} assert {variable.entity for variable in model.variables} <= entity_keys + + +# Read from the country packages pinned in pyproject.toml, independently of +# how the model version derives them. +PINNED_ENTITIES = { + "uk": [ + ("person", "people", True), + ("benunit", "benunits", False), + ("household", "households", False), + ], + "us": [ + ("household", "households", False), + ("spm_unit", "spm_units", False), + ("family", "families", False), + ("tax_unit", "tax_units", False), + ("marital_unit", "marital_units", False), + ("person", "people", True), + ], +} + + +def test_country_models_publish_their_pinned_entities(country): + assert [ + (entity.key, entity.plural, entity.is_person) + for entity in country.model.entities + ] == PINNED_ENTITIES[country.id] + assert {"employment_income", "age"} <= set(country.model.basic_inputs) diff --git a/tests/test_parameter_metadata.py b/tests/test_parameter_metadata.py index 6d07f50a..becfa2b2 100644 --- a/tests/test_parameter_metadata.py +++ b/tests/test_parameter_metadata.py @@ -110,3 +110,31 @@ def test_country_models_publish_each_parameters_flags(country): ), core_parameter.name assert any(not parameter.economy for parameter in model.parameters) assert any(not parameter.household for parameter in model.parameters) + + +# (economy, household) as the country packages' parameter files declare them, +# directly or through propagate_metadata_to_children. +PINNED_PARAMETER_FLAGS = { + "uk": { + "gov.dwp.ESA.eligibility.min_age": (False, True), + "gov.hmrc.minimum_wage.apprentice": (False, False), + "gov.hmrc.income_tax.rates.uk[0].rate": (True, True), + }, + "us": { + "gov.bls.cpi.cpi_u": (False, True), + "gov.aca.slcsp.max_adult_age": (False, False), + "gov.irs.income.bracket.rates.1": (True, True), + }, +} + + +def test_country_models_publish_pinned_parameter_flags(country): + published_flags = { + name: ( + country.model.get_parameter(name).economy, + country.model.get_parameter(name).household, + ) + for name in PINNED_PARAMETER_FLAGS[country.id] + } + + assert published_flags == PINNED_PARAMETER_FLAGS[country.id] diff --git a/tests/test_parameter_nodes.py b/tests/test_parameter_nodes.py index a0060ebc..552fd8e1 100644 --- a/tests/test_parameter_nodes.py +++ b/tests/test_parameter_nodes.py @@ -166,5 +166,23 @@ def test_country_models_publish_every_bracket_schedule(model, country_system_mod def test_uk_income_tax_schedule_is_a_labelled_node(): schedule = uk_latest.get_parameter_node("gov.hmrc.income_tax.rates.uk") - assert schedule.label + assert schedule.label == "Income Tax main rates" assert "gov.hmrc.income_tax.rates.uk[0]" in uk_latest.parameter_nodes_by_name + + +@pytest.mark.parametrize( + ("model", "schedule_name", "stored_label"), + [ + (uk_latest, "gov.hmrc.income_tax.rates.uk", "Income Tax main rates"), + (uk_latest, "gov.hmrc.income_tax.rates.dividends", None), + (us_latest, "gov.irs.income.bracket.rates", "Individual income tax rates"), + (us_latest, "gov.contrib.ubi_center.basic_income.amount.person.by_age", None), + ], + ids=["uk-labelled", "uk-unlabelled", "us-labelled", "us-unlabelled"], +) +def test_country_schedules_publish_their_stored_label_or_none( + model, schedule_name, stored_label +): + # A schedule without a stored label publishes None; consumers generate + # one from its name. + assert model.get_parameter_node(schedule_name).label == stored_label diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index e9e2d094..b411d83f 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -536,3 +536,43 @@ def test_us_model_publishes_parameter_named_parts(): == "gov.household.household_state_benefits" ) assert isinstance(household_state_benefits.adds, list) + + +def test_uk_model_publishes_pinned_variable_flags_and_parts(): + assert uk_latest.get_variable("age").is_input_variable is True + assert uk_latest.get_variable("employment_income_before_lsr").is_input_variable + assert uk_latest.get_variable("income_tax").is_input_variable is False + assert uk_latest.get_variable("employment_income").hidden_input is False + assert ( + uk_latest.get_variable("pension_contributions").adds_parameter_path + == "gov.hmrc.pensions.pensions_programs" + ) + assert ( + sum(1 for variable in uk_latest.variables if variable.adds_parameter_path) == 4 + ) + assert not any( + variable.subtracts_parameter_path for variable in uk_latest.variables + ) + + +def test_us_model_publishes_pinned_variable_flags_modules_and_parts(): + age = us_latest.get_variable("age") + income_tax = us_latest.get_variable("income_tax") + assert (age.is_input_variable, income_tax.is_input_variable) == (True, False) + assert age.module_name == "household.demographic.age.age" + assert income_tax.module_name == "gov.irs.tax.federal_income.income_tax" + assert us_latest.get_variable("medicaid_rating_area").hidden_input is True + assert us_latest.get_variable("employment_income").hidden_input is False + al_agi = us_latest.get_variable("al_agi") + assert al_agi.adds_parameter_path == ( + "gov.states.al.tax.income.agi.gross_income_sources" + ) + assert al_agi.subtracts_parameter_path == "gov.states.al.tax.income.agi.deductions" + assert ( + sum(1 for variable in us_latest.variables if variable.adds_parameter_path) + == 237 + ) + assert ( + sum(1 for variable in us_latest.variables if variable.subtracts_parameter_path) + == 2 + ) From f7124ab949ef2a152786e930414469dfb4b5bb83 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:31:41 +0400 Subject: [PATCH 11/14] Keep enum labels declared as numbers, booleans or None EnumValue.label held only a string or an integer, so a member declared as True or 1.0 became 1, and one declared as 0.5 or None failed validation, which would fail building the country model version. The label now holds any of those scalars and keeps each declared type, also through JSON. Co-Authored-By: Claude Opus 5.5 --- src/policyengine/core/variable.py | 6 +++-- tests/test_variable_metadata.py | 42 +++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/src/policyengine/core/variable.py b/src/policyengine/core/variable.py index 79a4bcef..3dc4951d 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -14,11 +14,13 @@ class EnumValue(BaseModel): value: str = Field( description="Name stored in household data, such as 'SINGLE'", ) - label: Union[str, int] = Field( + # Pydantic's smart union keeps each declared scalar's own type: True stays + # True, 1.0 stays 1.0 and None stays None, rather than being converted. + label: Union[str, int, float, bool, None] = Field( description=( "Label shown to users, such as 'Single': the value the country " "model declares for the member, unchanged, so a member declared " - "with an integer has an integer label" + "with an integer, number, boolean or None keeps that value" ), ) diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index b411d83f..96a7b96a 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -273,6 +273,16 @@ class _NumberedStatus(Enum): SECOND = 2 +class _MixedScalarStatus(Enum): + # Distinct values, since equal values (True, 1 and 1.0) would be aliases. + NAMED = "Named" + FLAGGED = True + COUNTED = 3 + WHOLE_RATE = 2.0 + PART_RATE = 0.5 + UNSET = None + + def _enum_core_variable(name: str, possible_values: type[Enum]) -> SimpleNamespace: first_member = next(iter(possible_values)) return SimpleNamespace( @@ -342,6 +352,38 @@ def test_enum_values_keep_integer_labels_as_declared(): } +def test_enum_values_keep_each_declared_scalar_type(): + model = _model_version() + mixed_status = _enum_core_variable("mixed_status", _MixedScalarStatus) + system = SimpleNamespace( + variables={mixed_status.name: mixed_status}, + parameters=None, + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + variable = model.get_variable("mixed_status") + published = [ + (enum_value.value, enum_value.label, type(enum_value.label)) + for enum_value in variable.enum_values + ] + assert published == [ + ("NAMED", "Named", str), + ("FLAGGED", True, bool), + ("COUNTED", 3, int), + ("WHOLE_RATE", 2.0, float), + ("PART_RATE", 0.5, float), + ("UNSET", None, type(None)), + ] + restored = [ + EnumValue.model_validate_json(enum_value.model_dump_json()) + for enum_value in variable.enum_values + ] + assert [(enum_value.label, type(enum_value.label)) for enum_value in restored] == [ + (label, label_type) for _, label, label_type in published + ] + + def test_variables_without_possible_values_have_no_enum_values(): model = _model_version() core_variable = SimpleNamespace( From f2ce69b96567e0871964d175b81c588eb1139fab Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:35:30 +0400 Subject: [PATCH 12/14] Share one skipping fixture for the new country-model tests The country tests this branch added now get the model version through shared fixtures (country_model for each country, uk_country_model and us_country_model), which skip when the country package is missing and import the model version only after that check. They and the empty model version helper are defined once, in tests/fixtures/country_model_fixtures.py, and registered in conftest. Tests that existed before this branch keep their module-level imports. Co-Authored-By: Claude Opus 5.5 --- tests/conftest.py | 5 + tests/fixtures/country_model_fixtures.py | 56 ++++++++++++ tests/test_model_entities.py | 56 +++--------- tests/test_parameter_metadata.py | 54 +++-------- tests/test_parameter_nodes.py | 69 +++++++------- tests/test_variable_metadata.py | 111 ++++++++++------------- 6 files changed, 167 insertions(+), 184 deletions(-) create mode 100644 tests/fixtures/country_model_fixtures.py diff --git a/tests/conftest.py b/tests/conftest.py index 0f730005..4969579e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,11 @@ import pytest # Import fixtures from fixtures module so pytest can discover them +from tests.fixtures.country_model_fixtures import ( # noqa: F401 + country_model, + uk_country_model, + us_country_model, +) from tests.fixtures.filtering_fixtures import ( # noqa: F401 create_uk_test_dataset, create_us_test_dataset, diff --git a/tests/fixtures/country_model_fixtures.py b/tests/fixtures/country_model_fixtures.py new file mode 100644 index 00000000..941638a1 --- /dev/null +++ b/tests/fixtures/country_model_fixtures.py @@ -0,0 +1,56 @@ +"""Country model versions for metadata tests, skipped when a package is missing. + +``policyengine.tax_benefit_models.uk`` and ``.us`` define ``uk_latest`` and +``us_latest`` only when the country package is installed, so these fixtures +check for the package first and import the model version afterwards. +""" + +import importlib +from types import SimpleNamespace + +import pytest + +from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion + + +def make_empty_model_version() -> TaxBenefitModelVersion: + """A model version with no country package, for populating from stubs.""" + return TaxBenefitModelVersion( + id="test@1", + model=TaxBenefitModel(id="test", name="Test"), + version="1", + ) + + +def load_country_model(country_id: str) -> SimpleNamespace: + """Return a country's published model version and its core system. + + The result has ``country_id``, ``model_version`` (``uk_latest`` or + ``us_latest``) and ``core_system`` (the country package's tax-benefit + system). Skips the calling test when the country package is not installed. + """ + core_system = pytest.importorskip(f"policyengine_{country_id}.system").system + country_module = importlib.import_module( + f"policyengine.tax_benefit_models.{country_id}" + ) + return SimpleNamespace( + country_id=country_id, + model_version=getattr(country_module, f"{country_id}_latest"), + core_system=core_system, + ) + + +@pytest.fixture(params=["uk", "us"]) +def country_model(request) -> SimpleNamespace: + """Each country's model version and core system, from load_country_model.""" + return load_country_model(request.param) + + +@pytest.fixture +def uk_country_model() -> SimpleNamespace: + return load_country_model("uk") + + +@pytest.fixture +def us_country_model() -> SimpleNamespace: + return load_country_model("us") diff --git a/tests/test_model_entities.py b/tests/test_model_entities.py index cba548f2..8f46456c 100644 --- a/tests/test_model_entities.py +++ b/tests/test_model_entities.py @@ -1,43 +1,13 @@ """Tests for the entities and basic inputs a model version publishes.""" -import importlib import json from types import SimpleNamespace -import pytest - -from policyengine.core import Entity, TaxBenefitModel, TaxBenefitModelVersion +from policyengine.core import Entity from policyengine.tax_benefit_models.common.model_version import ( MicrosimulationModelVersion, ) - - -def _model_version() -> TaxBenefitModelVersion: - return TaxBenefitModelVersion( - id="test@1", - model=TaxBenefitModel(id="test", name="Test"), - version="1", - ) - - -@pytest.fixture(params=["uk", "us"]) -def country(request) -> SimpleNamespace: - """The country's published model version and its core tax-benefit system. - - Skips when the country package is not installed. The model is imported - here, after that check, because the country module defines it only when - the package is installed. - """ - country_id = request.param - country_system = pytest.importorskip(f"policyengine_{country_id}.system").system - country_module = importlib.import_module( - f"policyengine.tax_benefit_models.{country_id}" - ) - return SimpleNamespace( - id=country_id, - model=getattr(country_module, f"{country_id}_latest"), - system=country_system, - ) +from tests.fixtures.country_model_fixtures import make_empty_model_version def _core_entity(key: str, plural: str, is_person: bool) -> SimpleNamespace: @@ -51,7 +21,7 @@ def _core_entity(key: str, plural: str, is_person: bool) -> SimpleNamespace: def test_populate_entities_records_key_plural_and_person_flag_in_order(): - model = _model_version() + model = make_empty_model_version() system = SimpleNamespace( entities=[ _core_entity("person", "people", True), @@ -72,7 +42,7 @@ def test_populate_entities_records_key_plural_and_person_flag_in_order(): def test_populate_entities_accepts_a_country_model_without_basic_inputs(): - model = _model_version() + model = make_empty_model_version() system = SimpleNamespace( entities=[_core_entity("person", "people", True)], basic_inputs=None, @@ -84,14 +54,14 @@ def test_populate_entities_accepts_a_country_model_without_basic_inputs(): def test_model_version_defaults_to_no_entities_or_basic_inputs(): - model = _model_version() + model = make_empty_model_version() assert model.entities == [] assert model.basic_inputs == [] def test_entities_and_basic_inputs_round_trip_through_json(): - model = _model_version() + model = make_empty_model_version() model.entities = [Entity(key="person", plural="people", is_person=True)] model.basic_inputs = ["age"] @@ -107,9 +77,9 @@ def test_entities_and_basic_inputs_round_trip_through_json(): assert restored_entities == model.entities -def test_country_models_publish_their_entities_and_basic_inputs(country): - model = country.model - country_system = country.system +def test_country_models_publish_their_entities_and_basic_inputs(country_model): + model = country_model.model_version + country_system = country_model.core_system assert [ (entity.key, entity.plural, entity.is_person) for entity in model.entities @@ -144,9 +114,9 @@ def test_country_models_publish_their_entities_and_basic_inputs(country): } -def test_country_models_publish_their_pinned_entities(country): +def test_country_models_publish_their_pinned_entities(country_model): assert [ (entity.key, entity.plural, entity.is_person) - for entity in country.model.entities - ] == PINNED_ENTITIES[country.id] - assert {"employment_income", "age"} <= set(country.model.basic_inputs) + for entity in country_model.model_version.entities + ] == PINNED_ENTITIES[country_model.country_id] + assert {"employment_income", "age"} <= set(country_model.model_version.basic_inputs) diff --git a/tests/test_parameter_metadata.py b/tests/test_parameter_metadata.py index becfa2b2..5c6fa6c9 100644 --- a/tests/test_parameter_metadata.py +++ b/tests/test_parameter_metadata.py @@ -1,45 +1,16 @@ """Tests for the economy and household flags each parameter publishes.""" -import importlib import json from types import SimpleNamespace -import pytest from policyengine_core.parameters import Parameter as CoreParameter from policyengine_core.parameters import ParameterNode as CoreParameterNode -from policyengine.core import Parameter, TaxBenefitModel, TaxBenefitModelVersion +from policyengine.core import Parameter from policyengine.tax_benefit_models.common.model_version import ( MicrosimulationModelVersion, ) - - -def _model_version() -> TaxBenefitModelVersion: - return TaxBenefitModelVersion( - id="test@1", - model=TaxBenefitModel(id="test", name="Test"), - version="1", - ) - - -@pytest.fixture(params=["uk", "us"]) -def country(request) -> SimpleNamespace: - """The country's published model version and its core tax-benefit system. - - Skips when the country package is not installed. The model is imported - here, after that check, because the country module defines it only when - the package is installed. - """ - country_id = request.param - country_system = pytest.importorskip(f"policyengine_{country_id}.system").system - country_module = importlib.import_module( - f"policyengine.tax_benefit_models.{country_id}" - ) - return SimpleNamespace( - id=country_id, - model=getattr(country_module, f"{country_id}_latest"), - system=country_system, - ) +from tests.fixtures.country_model_fixtures import make_empty_model_version def _flagged_parameter(metadata: dict) -> dict: @@ -51,7 +22,7 @@ def _flagged_parameter(metadata: dict) -> dict: def test_populate_parameters_records_economy_and_household_flags(): - model = _model_version() + model = make_empty_model_version() parameters = CoreParameterNode( "", data={ @@ -85,7 +56,7 @@ def test_populate_parameters_records_economy_and_household_flags(): def test_parameter_flags_default_to_offered(): - model = _model_version() + model = make_empty_model_version() manually_built = Parameter( name="gov.example.amount", tax_benefit_model_version=model ) @@ -94,9 +65,9 @@ def test_parameter_flags_default_to_offered(): assert manually_built.household is True -def test_country_models_publish_each_parameters_flags(country): - model = country.model - country_system = country.system +def test_country_models_publish_each_parameters_flags(country_model): + model = country_model.model_version + country_system = country_model.core_system for core_parameter in country_system.parameters.get_descendants(): if not isinstance(core_parameter, CoreParameter): @@ -128,13 +99,14 @@ def test_country_models_publish_each_parameters_flags(country): } -def test_country_models_publish_pinned_parameter_flags(country): +def test_country_models_publish_pinned_parameter_flags(country_model): + model = country_model.model_version published_flags = { name: ( - country.model.get_parameter(name).economy, - country.model.get_parameter(name).household, + model.get_parameter(name).economy, + model.get_parameter(name).household, ) - for name in PINNED_PARAMETER_FLAGS[country.id] + for name in PINNED_PARAMETER_FLAGS[country_model.country_id] } - assert published_flags == PINNED_PARAMETER_FLAGS[country.id] + assert published_flags == PINNED_PARAMETER_FLAGS[country_model.country_id] diff --git a/tests/test_parameter_nodes.py b/tests/test_parameter_nodes.py index 552fd8e1..3754a9d3 100644 --- a/tests/test_parameter_nodes.py +++ b/tests/test_parameter_nodes.py @@ -1,26 +1,19 @@ """Tests for public parameter-node projection from country models.""" -import importlib from types import SimpleNamespace import pytest from policyengine_core.parameters import ParameterNode as CoreParameterNode from policyengine_core.parameters import ParameterScale as CoreParameterScale -from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion from policyengine.tax_benefit_models.common.model_version import ( MicrosimulationModelVersion, ) from policyengine.tax_benefit_models.uk import uk_latest from policyengine.tax_benefit_models.us import us_latest - - -def _model_version() -> TaxBenefitModelVersion: - return TaxBenefitModelVersion( - id="test@1", - model=TaxBenefitModel(id="test", name="Test"), - version="1", - ) +from tests.fixtures.country_model_fixtures import ( + make_empty_model_version as _model_version, +) def test_populate_parameters_excludes_empty_organizational_nodes(): @@ -143,13 +136,9 @@ def test_populate_parameters_leaves_out_bracket_schedules_without_brackets(): assert "gov.tax.rates" in model.parameter_nodes_by_name -@pytest.mark.parametrize( - ("model", "country_system_module"), - [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], - ids=["uk", "us"], -) -def test_country_models_publish_every_bracket_schedule(model, country_system_module): - country_system = importlib.import_module(country_system_module).system +def test_country_models_publish_every_bracket_schedule(country_model): + model = country_model.model_version + country_system = country_model.core_system schedules = [ core_parameter @@ -163,26 +152,32 @@ def test_country_models_publish_every_bracket_schedule(model, country_system_mod assert schedule.description == core_schedule.description -def test_uk_income_tax_schedule_is_a_labelled_node(): - schedule = uk_latest.get_parameter_node("gov.hmrc.income_tax.rates.uk") +def test_uk_income_tax_schedule_is_a_node_above_its_brackets(uk_country_model): + model = uk_country_model.model_version - assert schedule.label == "Income Tax main rates" - assert "gov.hmrc.income_tax.rates.uk[0]" in uk_latest.parameter_nodes_by_name + assert "gov.hmrc.income_tax.rates.uk" in model.parameter_nodes_by_name + assert "gov.hmrc.income_tax.rates.uk[0]" in model.parameter_nodes_by_name -@pytest.mark.parametrize( - ("model", "schedule_name", "stored_label"), - [ - (uk_latest, "gov.hmrc.income_tax.rates.uk", "Income Tax main rates"), - (uk_latest, "gov.hmrc.income_tax.rates.dividends", None), - (us_latest, "gov.irs.income.bracket.rates", "Individual income tax rates"), - (us_latest, "gov.contrib.ubi_center.basic_income.amount.person.by_age", None), - ], - ids=["uk-labelled", "uk-unlabelled", "us-labelled", "us-unlabelled"], -) -def test_country_schedules_publish_their_stored_label_or_none( - model, schedule_name, stored_label -): - # A schedule without a stored label publishes None; consumers generate - # one from its name. - assert model.get_parameter_node(schedule_name).label == stored_label +# Stored labels in the country packages pinned in pyproject.toml. A schedule +# without one publishes None; consumers generate a label from its name. +PINNED_SCHEDULE_LABELS = { + "uk": { + "gov.hmrc.income_tax.rates.uk": "Income Tax main rates", + "gov.hmrc.income_tax.rates.dividends": None, + }, + "us": { + "gov.irs.income.bracket.rates": "Individual income tax rates", + "gov.contrib.ubi_center.basic_income.amount.person.by_age": None, + }, +} + + +def test_country_schedules_publish_their_stored_label_or_none(country_model): + model = country_model.model_version + pinned_labels = PINNED_SCHEDULE_LABELS[country_model.country_id] + + assert { + schedule_name: model.get_parameter_node(schedule_name).label + for schedule_name in pinned_labels + } == pinned_labels diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index 96a7b96a..0fb9fc82 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -1,20 +1,21 @@ """Tests for public variable metadata copied from country models.""" import datetime -import importlib import json from enum import Enum from types import SimpleNamespace -import pytest from policyengine_core.parameters import ParameterNode as CoreParameterNode -from policyengine.core import EnumValue, TaxBenefitModel, TaxBenefitModelVersion +from policyengine.core import EnumValue from policyengine.tax_benefit_models.common.model_version import ( MicrosimulationModelVersion, ) from policyengine.tax_benefit_models.uk import uk_latest from policyengine.tax_benefit_models.us import us_latest +from tests.fixtures.country_model_fixtures import ( + make_empty_model_version as _model_version, +) SELECTED_METADATA_FIELDS = { "definition_period", @@ -29,14 +30,6 @@ } -def _model_version() -> TaxBenefitModelVersion: - return TaxBenefitModelVersion( - id="test@1", - model=TaxBenefitModel(id="test", name="Test"), - version="1", - ) - - def test_populate_variables_exposes_selected_country_metadata(): model = _model_version() core_variable = SimpleNamespace( @@ -235,15 +228,11 @@ def test_input_flags_and_module_default_to_unrecorded_for_older_country_models() assert variable.module_name is None -@pytest.mark.parametrize( - ("model", "country_system_module"), - [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], - ids=["uk", "us"], -) def test_country_models_publish_input_flags_and_modules_as_core_declares_them( - model, country_system_module + country_model, ): - country_system = importlib.import_module(country_system_module).system + model = country_model.model_version + country_system = country_model.core_system for name, core_variable in country_system.variables.items(): variable = model.get_variable(name) @@ -254,9 +243,11 @@ def test_country_models_publish_input_flags_and_modules_as_core_declares_them( assert any(variable.is_input_variable is False for variable in model.variables) -def test_us_model_publishes_hidden_inputs(): +def test_us_model_publishes_hidden_inputs(us_country_model): hidden_inputs = [ - variable.name for variable in us_latest.variables if variable.hidden_input + variable.name + for variable in us_country_model.model_version.variables + if variable.hidden_input ] assert hidden_inputs, "The US package hides some inputs from input lists" @@ -406,13 +397,9 @@ def test_variables_without_possible_values_have_no_enum_values(): assert model.get_variable("minimal_input").enum_values is None -@pytest.mark.parametrize( - ("model", "country_system_module"), - [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], - ids=["uk", "us"], -) -def test_country_models_publish_every_enum_label(model, country_system_module): - country_system = importlib.import_module(country_system_module).system +def test_country_models_publish_every_enum_label(country_model): + model = country_model.model_version + country_system = country_model.core_system enum_variable_count = 0 for name, core_variable in country_system.variables.items(): @@ -434,15 +421,17 @@ def test_country_models_publish_every_enum_label(model, country_system_module): assert enum_variable_count > 0 -def test_us_filing_status_publishes_its_labels(): - filing_status = us_latest.get_variable("filing_status") +def test_us_filing_status_publishes_its_labels(us_country_model): + filing_status = us_country_model.model_version.get_variable("filing_status") assert EnumValue(value="SINGLE", label="Single") in filing_status.enum_values assert EnumValue(value="JOINT", label="Joint") in filing_status.enum_values -def test_us_integer_valued_enum_publishes_integer_labels(): - eligibility = us_latest.get_variable("il_personal_exemption_eligibility_status") +def test_us_integer_valued_enum_publishes_integer_labels(us_country_model): + eligibility = us_country_model.model_version.get_variable( + "il_personal_exemption_eligibility_status" + ) assert [ (enum_value.value, enum_value.label) for enum_value in eligibility.enum_values @@ -538,15 +527,11 @@ def test_parameter_path_is_published_when_the_parameter_cannot_be_resolved(): assert variable.adds_parameter_path == "gov.example.missing" -@pytest.mark.parametrize( - ("model", "country_system_module"), - [(uk_latest, "policyengine_uk.system"), (us_latest, "policyengine_us.system")], - ids=["uk", "us"], -) def test_country_models_publish_every_parameter_named_by_adds_or_subtracts( - model, country_system_module + country_model, ): - country_system = importlib.import_module(country_system_module).system + model = country_model.model_version + country_system = country_model.core_system for name, core_variable in country_system.variables.items(): variable = model.get_variable(name) @@ -560,19 +545,18 @@ def test_country_models_publish_every_parameter_named_by_adds_or_subtracts( assert published_path is None, (name, attribute) -def test_us_model_publishes_parameter_named_parts(): +def test_us_model_publishes_parameter_named_parts(us_country_model): + us_model = us_country_model.model_version adds_paths = [ - variable for variable in us_latest.variables if variable.adds_parameter_path + variable for variable in us_model.variables if variable.adds_parameter_path ] subtracts_paths = [ - variable - for variable in us_latest.variables - if variable.subtracts_parameter_path + variable for variable in us_model.variables if variable.subtracts_parameter_path ] assert len(adds_paths) >= 200 assert len(subtracts_paths) >= 1 - household_state_benefits = us_latest.get_variable("household_state_benefits") + household_state_benefits = us_model.get_variable("household_state_benefits") assert ( household_state_benefits.adds_parameter_path == "gov.household.household_state_benefits" @@ -580,41 +564,42 @@ def test_us_model_publishes_parameter_named_parts(): assert isinstance(household_state_benefits.adds, list) -def test_uk_model_publishes_pinned_variable_flags_and_parts(): - assert uk_latest.get_variable("age").is_input_variable is True - assert uk_latest.get_variable("employment_income_before_lsr").is_input_variable - assert uk_latest.get_variable("income_tax").is_input_variable is False - assert uk_latest.get_variable("employment_income").hidden_input is False +def test_uk_model_publishes_pinned_variable_flags_and_parts(uk_country_model): + uk_model = uk_country_model.model_version + assert uk_model.get_variable("age").is_input_variable is True + assert uk_model.get_variable("employment_income_before_lsr").is_input_variable + assert uk_model.get_variable("income_tax").is_input_variable is False + assert uk_model.get_variable("employment_income").hidden_input is False assert ( - uk_latest.get_variable("pension_contributions").adds_parameter_path + uk_model.get_variable("pension_contributions").adds_parameter_path == "gov.hmrc.pensions.pensions_programs" ) assert ( - sum(1 for variable in uk_latest.variables if variable.adds_parameter_path) == 4 - ) - assert not any( - variable.subtracts_parameter_path for variable in uk_latest.variables + sum(1 for variable in uk_model.variables if variable.adds_parameter_path) == 4 ) + assert not any(variable.subtracts_parameter_path for variable in uk_model.variables) -def test_us_model_publishes_pinned_variable_flags_modules_and_parts(): - age = us_latest.get_variable("age") - income_tax = us_latest.get_variable("income_tax") +def test_us_model_publishes_pinned_variable_flags_modules_and_parts( + us_country_model, +): + us_model = us_country_model.model_version + age = us_model.get_variable("age") + income_tax = us_model.get_variable("income_tax") assert (age.is_input_variable, income_tax.is_input_variable) == (True, False) assert age.module_name == "household.demographic.age.age" assert income_tax.module_name == "gov.irs.tax.federal_income.income_tax" - assert us_latest.get_variable("medicaid_rating_area").hidden_input is True - assert us_latest.get_variable("employment_income").hidden_input is False - al_agi = us_latest.get_variable("al_agi") + assert us_model.get_variable("medicaid_rating_area").hidden_input is True + assert us_model.get_variable("employment_income").hidden_input is False + al_agi = us_model.get_variable("al_agi") assert al_agi.adds_parameter_path == ( "gov.states.al.tax.income.agi.gross_income_sources" ) assert al_agi.subtracts_parameter_path == "gov.states.al.tax.income.agi.deductions" assert ( - sum(1 for variable in us_latest.variables if variable.adds_parameter_path) - == 237 + sum(1 for variable in us_model.variables if variable.adds_parameter_path) == 237 ) assert ( - sum(1 for variable in us_latest.variables if variable.subtracts_parameter_path) + sum(1 for variable in us_model.variables if variable.subtracts_parameter_path) == 2 ) From 9faae9b9b7acb840b3cc01220682ae17fb062202 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:37:10 +0400 Subject: [PATCH 13/14] Check parameter-named parts with lower bounds that allow additions The exact totals of parameter-named adds and subtracts would fail on a routine country package update that adds one, although .py's behavior is unchanged. The pinned variables stay; the counts are now at least the pinned packages' totals (UK 4 adds; US 237 adds and 2 subtracts), and every published path must name a published parameter. Co-Authored-By: Claude Opus 5.5 --- tests/test_variable_metadata.py | 45 +++++++++++++++------------------ 1 file changed, 20 insertions(+), 25 deletions(-) diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py index 0fb9fc82..02491377 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -545,23 +545,15 @@ def test_country_models_publish_every_parameter_named_by_adds_or_subtracts( assert published_path is None, (name, attribute) -def test_us_model_publishes_parameter_named_parts(us_country_model): - us_model = us_country_model.model_version - adds_paths = [ - variable for variable in us_model.variables if variable.adds_parameter_path - ] - subtracts_paths = [ - variable for variable in us_model.variables if variable.subtracts_parameter_path +def _published_parameter_paths(model, attribute: str) -> list[str]: + """Each published adds or subtracts parameter path, checking it resolves.""" + paths = [ + getattr(variable, f"{attribute}_parameter_path") + for variable in model.variables + if getattr(variable, f"{attribute}_parameter_path") ] - - assert len(adds_paths) >= 200 - assert len(subtracts_paths) >= 1 - household_state_benefits = us_model.get_variable("household_state_benefits") - assert ( - household_state_benefits.adds_parameter_path - == "gov.household.household_state_benefits" - ) - assert isinstance(household_state_benefits.adds, list) + assert all(path in model.parameters_by_name for path in paths), attribute + return paths def test_uk_model_publishes_pinned_variable_flags_and_parts(uk_country_model): @@ -574,10 +566,10 @@ def test_uk_model_publishes_pinned_variable_flags_and_parts(uk_country_model): uk_model.get_variable("pension_contributions").adds_parameter_path == "gov.hmrc.pensions.pensions_programs" ) - assert ( - sum(1 for variable in uk_model.variables if variable.adds_parameter_path) == 4 - ) - assert not any(variable.subtracts_parameter_path for variable in uk_model.variables) + # At least as many as the pinned UK package declares, so additions in a + # later package do not fail the test. + assert len(_published_parameter_paths(uk_model, "adds")) >= 4 + _published_parameter_paths(uk_model, "subtracts") def test_us_model_publishes_pinned_variable_flags_modules_and_parts( @@ -596,10 +588,13 @@ def test_us_model_publishes_pinned_variable_flags_modules_and_parts( "gov.states.al.tax.income.agi.gross_income_sources" ) assert al_agi.subtracts_parameter_path == "gov.states.al.tax.income.agi.deductions" + household_state_benefits = us_model.get_variable("household_state_benefits") assert ( - sum(1 for variable in us_model.variables if variable.adds_parameter_path) == 237 - ) - assert ( - sum(1 for variable in us_model.variables if variable.subtracts_parameter_path) - == 2 + household_state_benefits.adds_parameter_path + == "gov.household.household_state_benefits" ) + assert isinstance(household_state_benefits.adds, list) + # At least as many as the pinned US package declares, so additions in a + # later package do not fail the test. + assert len(_published_parameter_paths(us_model, "adds")) >= 237 + assert len(_published_parameter_paths(us_model, "subtracts")) >= 2 From b6a719fb5c13190c597320993400c4b717b87517 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:37:56 +0400 Subject: [PATCH 14/14] Add the changelog fragment for the published model metadata Co-Authored-By: Claude Opus 5.5 --- changelog.d/700.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/700.added.md diff --git a/changelog.d/700.added.md b/changelog.d/700.added.md new file mode 100644 index 00000000..67131ef3 --- /dev/null +++ b/changelog.d/700.added.md @@ -0,0 +1 @@ +Model versions publish the metadata the Public API could only read through the private `_load_system()` method: each model version's `entities` (key, plural and person flag) and `basic_inputs`; each variable's `is_input_variable`, `hidden_input`, `module_name`, `enum_values` (each value with its label, as declared) and, where its `adds` or `subtracts` names a parameter, `adds_parameter_path` or `subtracts_parameter_path`; each parameter's `economy` and `household` flags; and each bracket schedule as a parameter node with its stored label. Existing fields are unchanged: `adds` and `subtracts` still hold the list as of 2025-01-01.