From 3de153bbcc4219964460cbdf52e1c5e10d97435f Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 6 Feb 2026 21:44:05 +0300 Subject: [PATCH 1/4] feat: Add default_value and value_type to Variable class MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add default_value (Any) and value_type (type) fields to Variable model - Update US and UK models to serialize default_value for JSON compatibility: - Enum values converted to their .name string - datetime.date values converted to ISO format string - Primitives (bool, int, float, str) kept as-is - value_type preserves the original type for downstream consumers Closes #226 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- src/policyengine/core/variable.py | 2 ++ src/policyengine/tax_benefit_models/uk/model.py | 9 +++++++++ src/policyengine/tax_benefit_models/us/model.py | 9 +++++++++ 3 files changed, 20 insertions(+) diff --git a/src/policyengine/core/variable.py b/src/policyengine/core/variable.py index 24375120..fa8e50e2 100644 --- a/src/policyengine/core/variable.py +++ b/src/policyengine/core/variable.py @@ -13,3 +13,5 @@ class Variable(BaseModel): description: str | None = None data_type: type = None possible_values: list[Any] | None = None + default_value: Any = None + value_type: type | None = None diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index f5eae9c5..fac5b91f 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -126,6 +126,13 @@ def __init__(self, **kwargs: dict): self.id = f"{self.model.id}@{self.version}" for var_obj in system.variables.values(): + # Serialize default_value for JSON compatibility + default_val = var_obj.default_value + if var_obj.value_type is Enum: + default_val = default_val.name + elif var_obj.value_type is datetime.date: + default_val = default_val.isoformat() + variable = Variable( id=self.id + "-" + var_obj.name, name=var_obj.name, @@ -135,6 +142,8 @@ def __init__(self, **kwargs: dict): data_type=var_obj.value_type if var_obj.value_type is not Enum else str, + default_value=default_val, + value_type=var_obj.value_type, ) if ( hasattr(var_obj, "possible_values") diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 939f3f17..487e4d51 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -119,6 +119,13 @@ def __init__(self, **kwargs: dict): self.id = f"{self.model.id}@{self.version}" for var_obj in system.variables.values(): + # Serialize default_value for JSON compatibility + default_val = var_obj.default_value + if var_obj.value_type is Enum: + default_val = default_val.name + elif var_obj.value_type is datetime.date: + default_val = default_val.isoformat() + variable = Variable( id=self.id + "-" + var_obj.name, name=var_obj.name, @@ -128,6 +135,8 @@ def __init__(self, **kwargs: dict): data_type=var_obj.value_type if var_obj.value_type is not Enum else str, + default_value=default_val, + value_type=var_obj.value_type, ) if ( hasattr(var_obj, "possible_values") From 21855242df573e36d925976f2cb5f59a8800bd56 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 7 Feb 2026 03:42:51 +0300 Subject: [PATCH 2/4] test: Add Variable default_value integration tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add tests verifying US age variable has default_value of 40 - Add tests verifying enum variables have string default_value - Add tests verifying variables have value_type set - Add tests verifying UK model variable default_value 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- tests/test_models.py | 38 +++++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index 06045df9..3132abdf 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,4 +1,4 @@ -"""Tests for UK and US tax-benefit model versions.""" +"""Tests for UK and US tax-benefit model versions and core models.""" import re @@ -146,3 +146,39 @@ def test__given_bracket_label__then_follows_expected_format(self): f"Label '{p.label}' doesn't match expected bracket format" ) break + + +class TestVariableDefaultValue: + """Tests for Variable default_value and value_type fields.""" + + def test_us_age_variable_has_default_value_40(self): + """US age variable should have default_value of 40.""" + age_var = next((v for v in us_latest.variables if v.name == "age"), None) + assert age_var is not None, "age variable not found in US model" + assert age_var.default_value == 40, ( + f"Expected age default_value to be 40, got {age_var.default_value}" + ) + + def test_us_enum_variable_has_string_default_value(self): + """Enum variables should have string default_value (not enum object).""" + # age_group is an enum with default WORKING_AGE + age_group_var = next( + (v for v in us_latest.variables if v.name == "age_group"), None + ) + assert age_group_var is not None, "age_group variable not found in US model" + assert age_group_var.default_value == "WORKING_AGE", ( + f"Expected age_group default_value to be 'WORKING_AGE', " + f"got {age_group_var.default_value}" + ) + + def test_us_variables_have_value_type(self): + """US variables should have value_type set.""" + age_var = next((v for v in us_latest.variables if v.name == "age"), None) + assert age_var is not None, "age variable not found in US model" + assert age_var.value_type is not None, "age variable should have value_type" + + def test_uk_age_variable_has_default_value(self): + """UK age variable should have default_value set.""" + age_var = next((v for v in uk_latest.variables if v.name == "age"), None) + assert age_var is not None, "age variable not found in UK model" + assert age_var.default_value is not None, "UK age should have default_value" From 7c18322e6adbfdf9db1f30a22dc2dffd96748bad Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 16 Feb 2026 20:44:33 +0100 Subject: [PATCH 3/4] ci: Remove base branch filter from PR workflow Allow PR checks to run for PRs targeting any branch, not just main. Co-Authored-By: Claude Opus 4.6 --- .github/workflows/pr_code_changes.yaml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/pr_code_changes.yaml b/.github/workflows/pr_code_changes.yaml index aa03bc02..3d1a5139 100644 --- a/.github/workflows/pr_code_changes.yaml +++ b/.github/workflows/pr_code_changes.yaml @@ -3,9 +3,6 @@ name: Code changes on: pull_request: - branches: - - main - paths: - src/** - tests/** From 546500884c23ab14a78c99c5638bd8095d63cb3e Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 16 Feb 2026 20:59:32 +0100 Subject: [PATCH 4/4] fix: Use StrEnum instead of (str, Enum) for ruff UP042 Co-Authored-By: Claude Opus 4.6 --- src/policyengine/outputs/aggregate.py | 4 ++-- src/policyengine/outputs/change_aggregate.py | 4 ++-- src/policyengine/outputs/poverty.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/policyengine/outputs/aggregate.py b/src/policyengine/outputs/aggregate.py index 2d41259c..cf4dbca4 100644 --- a/src/policyengine/outputs/aggregate.py +++ b/src/policyengine/outputs/aggregate.py @@ -1,10 +1,10 @@ -from enum import Enum +from enum import StrEnum from typing import Any from policyengine.core import Output, Simulation -class AggregateType(str, Enum): +class AggregateType(StrEnum): SUM = "sum" MEAN = "mean" COUNT = "count" diff --git a/src/policyengine/outputs/change_aggregate.py b/src/policyengine/outputs/change_aggregate.py index b5bfe2df..29fe5069 100644 --- a/src/policyengine/outputs/change_aggregate.py +++ b/src/policyengine/outputs/change_aggregate.py @@ -1,10 +1,10 @@ -from enum import Enum +from enum import StrEnum from typing import Any from policyengine.core import Output, Simulation -class ChangeAggregateType(str, Enum): +class ChangeAggregateType(StrEnum): COUNT = "count" SUM = "sum" MEAN = "mean" diff --git a/src/policyengine/outputs/poverty.py b/src/policyengine/outputs/poverty.py index 9b5074f4..4955b814 100644 --- a/src/policyengine/outputs/poverty.py +++ b/src/policyengine/outputs/poverty.py @@ -1,6 +1,6 @@ """Poverty analysis output types.""" -from enum import Enum +from enum import StrEnum from typing import Any import pandas as pd @@ -9,7 +9,7 @@ from policyengine.core import Output, OutputCollection, Simulation -class UKPovertyType(str, Enum): +class UKPovertyType(StrEnum): """UK poverty measure types.""" ABSOLUTE_BHC = "absolute_bhc" @@ -18,7 +18,7 @@ class UKPovertyType(str, Enum): RELATIVE_AHC = "relative_ahc" -class USPovertyType(str, Enum): +class USPovertyType(StrEnum): """US poverty measure types.""" SPM = "spm"