From 4f29c968ea9075b99115afa79d840b7e5c9955e1 Mon Sep 17 00:00:00 2001 From: SakshiKekre Date: Tue, 3 Mar 2026 20:08:40 +0530 Subject: [PATCH] fix: Populate possible_values and set data_type='Enum' for enum variables MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a variable has possible_values (from policyengine), populate the possible_values field with JSON-encoded values and set data_type to "Enum". This allows the frontend to render dropdowns for enum variables like state_name. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- scripts/seed_models.py | 27 +++++++++++++++++++++---- src/policyengine_api/models/variable.py | 2 +- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/scripts/seed_models.py b/scripts/seed_models.py index 049558e..9313faa 100644 --- a/scripts/seed_models.py +++ b/scripts/seed_models.py @@ -28,6 +28,26 @@ ) +def _get_variable_type_info(var) -> tuple[str, str | None]: + """Extract data_type and possible_values from a policyengine variable. + + For enum variables (those with possible_values), returns ("Enum", json_values). + For other variables, returns (python_type_name, None). + + Returns: + Tuple of (data_type, possible_values_json) + """ + if var.possible_values: + return "Enum", json.dumps(var.possible_values) + + data_type = ( + var.data_type.__name__ + if hasattr(var.data_type, "__name__") + else str(var.data_type) + ) + return data_type, None + + def seed_model( model_version, session: Session, @@ -121,16 +141,15 @@ def seed_model( total=len(variables), ) for var in variables: + data_type, possible_values = _get_variable_type_info(var) var_rows.append( { "id": uuid4(), "name": var.name, "entity": var.entity, "description": var.description or "", - "data_type": var.data_type.__name__ - if hasattr(var.data_type, "__name__") - else str(var.data_type), - "possible_values": None, + "data_type": data_type, + "possible_values": possible_values, "tax_benefit_model_version_id": db_version.id, "created_at": datetime.now(timezone.utc), } diff --git a/src/policyengine_api/models/variable.py b/src/policyengine_api/models/variable.py index eeebddc..b147bcf 100644 --- a/src/policyengine_api/models/variable.py +++ b/src/policyengine_api/models/variable.py @@ -15,7 +15,7 @@ class VariableBase(SQLModel): entity: str description: str | None = None data_type: str | None = None # Store as string representation - possible_values: str | None = Field( + possible_values: list[str] | None = Field( default=None, sa_column=Column(JSON) ) # Store as JSON list default_value: Any = Field(