diff --git a/changelog.d/504.added.md b/changelog.d/504.added.md new file mode 100644 index 00000000..971b8c02 --- /dev/null +++ b/changelog.d/504.added.md @@ -0,0 +1 @@ +Expose variable period, unit, quantity type, references, applicability, numeric bounds, period-size behavior, and sanitized country metadata on the public `Variable` model. Exclude empty organizational nodes from the public parameter catalogue. diff --git a/src/policyengine/core/variable.py b/src/policyengine/core/variable.py index 03e53495..415081a7 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -1,6 +1,9 @@ -from typing import Any, Optional +import datetime +from collections.abc import Mapping +from enum import Enum +from typing import Any, Optional, Union -from pydantic import BaseModel +from pydantic import BaseModel, Field, field_validator from .tax_benefit_model_version import TaxBenefitModelVersion @@ -18,3 +21,51 @@ class Variable(BaseModel): value_type: Optional[type] = None adds: Optional[list[str]] = None subtracts: Optional[list[str]] = None + definition_period: Optional[str] = None + unit: Optional[str] = None + quantity_type: Optional[str] = None + reference: Optional[list[Any]] = None + defined_for: Optional[str] = None + min_value: Optional[Union[int, float]] = None + max_value: Optional[Union[int, float]] = None + is_period_size_independent: Optional[bool] = None + metadata: dict[str, Any] = Field(default_factory=dict) + + @field_validator("reference", mode="before") + @classmethod + def normalize_reference(cls, value: Any) -> Optional[list[Any]]: + """Return references as a JSON-safe list.""" + + if value is None: + return None + references = value if isinstance(value, (list, tuple)) else [value] + return [_json_safe(item) for item in references] + + @field_validator("metadata", mode="before") + @classmethod + def normalize_metadata(cls, value: Any) -> dict[str, Any]: + """Return arbitrary country metadata as a JSON-safe mapping.""" + + if value is None: + return {} + if not isinstance(value, Mapping): + raise TypeError("Variable metadata must be a mapping") + return {str(key): _json_safe(item) for key, item in value.items()} + + +def _json_safe(value: Any) -> Any: + """Normalize country-model metadata without retaining runtime objects.""" + + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, (datetime.date, datetime.datetime)): + return value.isoformat() + if isinstance(value, Enum): + return _json_safe(value.value) + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + return [_json_safe(item) for item in value] + if hasattr(value, "model_dump"): + return _json_safe(value.model_dump(mode="json")) + return str(value) diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index 57ac573b..c5e6c090 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -206,6 +206,17 @@ def _populate_variables(self, system) -> None: ), default_value=default_val, value_type=var_obj.value_type, + definition_period=getattr(var_obj, "definition_period", None), + unit=getattr(var_obj, "unit", None), + quantity_type=getattr(var_obj, "quantity_type", None), + reference=getattr(var_obj, "reference", None), + defined_for=getattr(var_obj, "defined_for", None), + min_value=getattr(var_obj, "min_value", None), + max_value=getattr(var_obj, "max_value", None), + is_period_size_independent=getattr( + var_obj, "is_period_size_independent", None + ), + metadata=getattr(var_obj, "metadata", None) or {}, ) if ( hasattr(var_obj, "possible_values") @@ -255,6 +266,8 @@ def _populate_parameters(self, system) -> None: ) self.add_parameter(parameter) elif isinstance(param_node, CoreParameterNode): + if not param_node.children: + continue node = ParameterNode( id=self.id + "-" + param_node.name, name=param_node.name, diff --git a/tests/test_parameter_nodes.py b/tests/test_parameter_nodes.py new file mode 100644 index 00000000..99a8edd6 --- /dev/null +++ b/tests/test_parameter_nodes.py @@ -0,0 +1,56 @@ +"""Tests for public parameter-node projection from country models.""" + +from types import SimpleNamespace + +import pytest +from policyengine_core.parameters import ParameterNode as CoreParameterNode + +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", + ) + + +def test_populate_parameters_excludes_empty_organizational_nodes(): + populated_node = CoreParameterNode( + "gov.example", + data={ + "amount": { + "description": "Example amount", + "values": {"2026-01-01": 10}, + } + }, + ) + empty_node = CoreParameterNode("generated_artifacts", data={}) + descendants = [populated_node, empty_node] + system = SimpleNamespace( + parameters=SimpleNamespace( + get_descendants=lambda: iter(descendants), + ) + ) + model = _model_version() + + MicrosimulationModelVersion._populate_parameters(model, system) + + assert [node.name for node in model.parameter_nodes] == ["gov.example"] + assert set(model.parameter_nodes_by_name) == {"gov.example"} + + +@pytest.mark.parametrize("model", [uk_latest, us_latest], ids=["uk", "us"]) +def test_country_models_do_not_expose_python_cache_nodes(model): + assert all( + "__pycache__" not in node.name.split(".") for node in model.parameter_nodes + ) + assert all( + "__pycache__" not in name.split(".") for name in model.parameter_nodes_by_name + ) diff --git a/tests/test_variable_metadata.py b/tests/test_variable_metadata.py new file mode 100644 index 00000000..575d77e5 --- /dev/null +++ b/tests/test_variable_metadata.py @@ -0,0 +1,144 @@ +"""Tests for public variable metadata copied from country models.""" + +import datetime +import json +from types import SimpleNamespace + +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 + +SELECTED_METADATA_FIELDS = { + "definition_period", + "unit", + "quantity_type", + "reference", + "defined_for", + "min_value", + "max_value", + "is_period_size_independent", + "metadata", +} + + +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( + name="sample_income", + label="Sample income", + entity=SimpleNamespace(key="person"), + documentation="Income used to test metadata projection.", + value_type=float, + default_value=0.0, + possible_values=None, + adds=None, + subtracts=None, + definition_period="year", + unit="currency-GBP", + quantity_type="flow", + reference=[ + "Example Act 2026", + {"title": "Example guidance", "published": datetime.date(2026, 1, 1)}, + ], + defined_for="is_adult", + min_value=0, + max_value=1_000_000.0, + is_period_size_independent=False, + metadata={"source": {"years": (2025, 2026)}}, + ) + system = SimpleNamespace( + variables={core_variable.name: core_variable}, + parameters=None, + ) + + MicrosimulationModelVersion._populate_variables(model, system) + + variable = model.get_variable("sample_income") + assert variable.definition_period == "year" + assert variable.unit == "currency-GBP" + assert variable.quantity_type == "flow" + assert variable.reference == [ + "Example Act 2026", + {"title": "Example guidance", "published": "2026-01-01"}, + ] + assert variable.defined_for == "is_adult" + assert variable.min_value == 0 + assert variable.max_value == 1_000_000.0 + assert variable.is_period_size_independent is False + assert variable.metadata == {"source": {"years": [2025, 2026]}} + + payload = json.loads(variable.model_dump_json(include=SELECTED_METADATA_FIELDS)) + assert payload["reference"][1]["published"] == "2026-01-01" + assert payload["metadata"] == {"source": {"years": [2025, 2026]}} + + +def test_selected_metadata_defaults_preserve_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.definition_period is None + assert variable.unit is None + assert variable.quantity_type is None + assert variable.reference is None + assert variable.defined_for is None + assert variable.min_value is None + assert variable.max_value is None + assert variable.is_period_size_independent is None + assert variable.metadata == {} + + +def test_uk_model_exposes_selected_variable_metadata(): + employment_income = uk_latest.get_variable("employment_income") + assert employment_income.definition_period == "year" + assert employment_income.unit == "currency-GBP" + assert employment_income.quantity_type == "flow" + assert employment_income.reference == [ + "Income Tax (Earnings and Pensions) Act 2003 s. 1(1)(a)" + ] + assert employment_income.is_period_size_independent is False + assert employment_income.metadata == {} + + personal_rent = uk_latest.get_variable("personal_rent") + assert personal_rent.defined_for == "is_household_head" + assert personal_rent.min_value is None + assert personal_rent.max_value is None + + +def test_us_model_exposes_selected_variable_metadata(): + employment_income = us_latest.get_variable("employment_income") + assert employment_income.definition_period == "year" + assert employment_income.unit == "currency-USD" + assert employment_income.quantity_type == "flow" + assert employment_income.reference == [ + "https://www.law.cornell.edu/uscode/text/26/3401#a" + ] + assert employment_income.is_period_size_independent is False + assert employment_income.metadata == {}