From 339f939d0e169b6d4d65b99cc99d2284c59927eb Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 6 Feb 2026 21:44:05 +0300 Subject: [PATCH] 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")