diff --git a/changelog.d/700.added.md b/changelog.d/700.added.md new file mode 100644 index 000000000..67131ef3e --- /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. diff --git a/src/policyengine/core/__init__.py b/src/policyengine/core/__init__.py index 703917460..b30b34fb9 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 @@ -31,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/entity.py b/src/policyengine/core/entity.py new file mode 100644 index 000000000..15375b610 --- /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/parameter.py b/src/policyengine/core/parameter.py index 680120cf4..5f61ba51a 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/core/parameter_node.py b/src/policyengine/core/parameter_node.py index 54d384a54..af060b799 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/core/tax_benefit_model_version.py b/src/policyengine/core/tax_benefit_model_version.py index 977561c1c..dc7a2de4e 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/core/variable.py b/src/policyengine/core/variable.py index 415081a7e..3dc4951d2 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -8,6 +8,23 @@ 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'", + ) + # 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, number, boolean or None keeps that value" + ), + ) + + class Variable(BaseModel): id: str name: str @@ -30,6 +47,50 @@ 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", + ) + 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=( + "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 3b72ad2c3..21c27b8d1 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -30,6 +30,8 @@ import pandas as pd from policyengine.core import ( + Entity, + EnumValue, Parameter, ParameterNode, StaleOutputError, @@ -93,6 +95,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. @@ -166,6 +176,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 +200,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 ( @@ -225,6 +248,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") @@ -236,13 +262,19 @@ def _populate_variables(self, system) -> None: var_obj.possible_values._value2member_map_.values(), ) ) + variable.enum_values = [ + EnumValue(value=member.name, label=member.value) + 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"))) @@ -255,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) @@ -270,9 +303,24 @@ 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) + 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/conftest.py b/tests/conftest.py index 0f7300054..4969579e2 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 000000000..941638a1c --- /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 new file mode 100644 index 000000000..8f46456ca --- /dev/null +++ b/tests/test_model_entities.py @@ -0,0 +1,122 @@ +"""Tests for the entities and basic inputs a model version publishes.""" + +import json +from types import SimpleNamespace + +from policyengine.core import Entity +from policyengine.tax_benefit_models.common.model_version import ( + MicrosimulationModelVersion, +) +from tests.fixtures.country_model_fixtures import make_empty_model_version + + +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 = make_empty_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 = make_empty_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 = make_empty_model_version() + + assert model.entities == [] + assert model.basic_inputs == [] + + +def test_entities_and_basic_inputs_round_trip_through_json(): + model = make_empty_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 + + +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 + ] == [ + (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 + + +# 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_model): + assert [ + (entity.key, entity.plural, entity.is_person) + 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 new file mode 100644 index 000000000..5c6fa6c9b --- /dev/null +++ b/tests/test_parameter_metadata.py @@ -0,0 +1,112 @@ +"""Tests for the economy and household flags each parameter publishes.""" + +import json +from types import SimpleNamespace + +from policyengine_core.parameters import Parameter as CoreParameter +from policyengine_core.parameters import ParameterNode as CoreParameterNode + +from policyengine.core import Parameter +from policyengine.tax_benefit_models.common.model_version import ( + MicrosimulationModelVersion, +) +from tests.fixtures.country_model_fixtures import make_empty_model_version + + +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 = make_empty_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 = make_empty_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 + + +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): + 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) + + +# (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_model): + model = country_model.model_version + published_flags = { + name: ( + model.get_parameter(name).economy, + model.get_parameter(name).household, + ) + for name in PINNED_PARAMETER_FLAGS[country_model.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 99a8edd6b..3754a9d3c 100644 --- a/tests/test_parameter_nodes.py +++ b/tests/test_parameter_nodes.py @@ -4,21 +4,16 @@ 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(): @@ -54,3 +49,135 @@ 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)" + ) + + +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 + + +def test_country_models_publish_every_bracket_schedule(country_model): + model = country_model.model_version + country_system = country_model.core_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_node_above_its_brackets(uk_country_model): + model = uk_country_model.model_version + + 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 + + +# 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 4bfa3d5a8..02491377a 100644 --- a/tests/test_variable_metadata.py +++ b/tests/test_variable_metadata.py @@ -2,14 +2,20 @@ import datetime import json +from enum import Enum from types import SimpleNamespace -from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion +from policyengine_core.parameters import ParameterNode as CoreParameterNode + +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", @@ -24,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( @@ -142,3 +140,461 @@ 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 + + +def test_country_models_publish_input_flags_and_modules_as_core_declares_them( + country_model, +): + 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) + 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(us_country_model): + hidden_inputs = [ + 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" + + +class _FilingStatus(Enum): + SINGLE = "Single" + JOINT = "Joint" + HEAD_OF_HOUSEHOLD = "Head of household" + + +class _NumberedStatus(Enum): + FIRST = 1 + 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( + 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) + system = SimpleNamespace( + variables={filing_status.name: filing_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"}, + ] + } + + +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_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( + 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 + + +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(): + 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, 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 + assert enum_variable_count > 0 + + +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(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 + ] == [("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( + "", + 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" + + +def test_country_models_publish_every_parameter_named_by_adds_or_subtracts( + country_model, +): + 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) + 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 _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 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): + 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_model.get_variable("pension_contributions").adds_parameter_path + == "gov.hmrc.pensions.pensions_programs" + ) + # 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( + 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_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" + 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) + # 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