diff --git a/src/policyengine/utils/parameter_labels.py b/src/policyengine/utils/parameter_labels.py index aac76e1..2fd3e25 100644 --- a/src/policyengine/utils/parameter_labels.py +++ b/src/policyengine/utils/parameter_labels.py @@ -26,34 +26,140 @@ def generate_label_for_parameter(param_node, system, scale_lookup): if "[" in param_name: return _generate_bracket_label(param_name, scale_lookup) - if param_node.parent and param_node.parent.metadata.get("breakdown"): - return _generate_breakdown_label(param_node, system) + # Check for breakdown - either direct child or nested + breakdown_parent = _find_breakdown_parent(param_node) + if breakdown_parent: + return _generate_breakdown_label(param_node, system, breakdown_parent) return None -def _generate_breakdown_label(param_node, system): - """Generate label for a breakdown parameter using enum values.""" - parent = param_node.parent - parent_label = parent.metadata.get("label") - breakdown_vars = parent.metadata.get("breakdown", []) +def _find_breakdown_parent(param_node): + """ + Walk up the tree to find the nearest ancestor with breakdown metadata. + + Args: + param_node: The CoreParameter object + + Returns: + The breakdown parent node, or None if not found + """ + current = param_node.parent + while current: + if current.metadata.get("breakdown"): + return current + current = getattr(current, "parent", None) + return None + + +def _generate_breakdown_label(param_node, system, breakdown_parent=None): + """ + Generate label for a breakdown parameter using enum values. + + Handles both single-level and nested breakdowns by walking up to the + breakdown parent and collecting all dimension values. + Args: + param_node: The CoreParameter object + system: The tax-benefit system + breakdown_parent: The ancestor node with breakdown metadata (optional) + + Returns: + str or None: Generated label, or None if cannot generate + """ + # Find breakdown parent if not provided + if breakdown_parent is None: + breakdown_parent = _find_breakdown_parent(param_node) + if not breakdown_parent: + return None + + parent_label = breakdown_parent.metadata.get("label") if not parent_label: return None - child_key = param_node.name.split(".")[-1] + breakdown_vars = breakdown_parent.metadata.get("breakdown", []) + breakdown_labels = breakdown_parent.metadata.get("breakdown_labels", []) + + # Collect dimension values from breakdown parent to param_node + dimension_values = _collect_dimension_values( + param_node, breakdown_parent + ) + + if not dimension_values: + return None + + # Generate labels for each dimension + formatted_parts = [] + for i, (dim_key, dim_value) in enumerate(dimension_values): + var_name = breakdown_vars[i] if i < len(breakdown_vars) else None + dim_label = breakdown_labels[i] if i < len(breakdown_labels) else None - for var_name in breakdown_vars: + formatted_value = _format_dimension_value( + dim_value, var_name, dim_label, system + ) + formatted_parts.append(formatted_value) + + return f"{parent_label} ({', '.join(formatted_parts)})" + + +def _collect_dimension_values(param_node, breakdown_parent): + """ + Collect dimension keys and values from breakdown parent to param_node. + + Args: + param_node: The CoreParameter object + breakdown_parent: The ancestor node with breakdown metadata + + Returns: + list of (dimension_key, value) tuples, ordered from parent to child + """ + # Build path from param_node up to breakdown_parent + path = [] + current = param_node + while current and current != breakdown_parent: + path.append(current) + current = getattr(current, "parent", None) + + # Reverse to get parent-to-child order + path.reverse() + + # Extract dimension values + dimension_values = [] + for i, node in enumerate(path): + key = node.name.split(".")[-1] + dimension_values.append((i, key)) + + return dimension_values + + +def _format_dimension_value(value, var_name, dim_label, system): + """ + Format a single dimension value with semantic label if available. + + Args: + value: The raw dimension value (e.g., "SINGLE", "1", "CA") + var_name: The breakdown variable name (e.g., "filing_status", "range(1, 9)") + dim_label: The human-readable label for this dimension (e.g., "Household size") + system: The tax-benefit system + + Returns: + str: Formatted dimension value + """ + # First, try to get enum display value + if var_name and isinstance(var_name, str) and not var_name.startswith("range(") and not var_name.startswith("list("): var = system.variables.get(var_name) if var and hasattr(var, "possible_values") and var.possible_values: - enum_class = var.possible_values try: - enum_value = enum_class[child_key].value - return f"{parent_label} ({enum_value})" + enum_value = var.possible_values[value].value + return str(enum_value) except (KeyError, AttributeError): - continue + pass + + # For range() dimensions or when no enum found, use breakdown_label if available + if dim_label: + return f"{dim_label} {value}" - return f"{parent_label} ({child_key})" + return value def _generate_bracket_label(param_name, scale_lookup): diff --git a/tests/fixtures/parameter_labels_fixtures.py b/tests/fixtures/parameter_labels_fixtures.py index e43263c..0e22424 100644 --- a/tests/fixtures/parameter_labels_fixtures.py +++ b/tests/fixtures/parameter_labels_fixtures.py @@ -11,6 +11,7 @@ class MockFilingStatus(Enum): SINGLE = "Single" JOINT = "Joint" HEAD_OF_HOUSEHOLD = "Head of household" + MARRIED_FILING_JOINTLY = "Married filing jointly" class MockStateCode(Enum): @@ -38,16 +39,21 @@ def create_mock_parent_node( name: str, label: str | None = None, breakdown: list[str] | None = None, + breakdown_labels: list[str] | None = None, + parent: Any = None, ) -> MagicMock: """Create a mock parent ParameterNode with optional breakdown metadata.""" - parent = MagicMock() - parent.name = name - parent.metadata = {} + node = MagicMock() + node.name = name + node.metadata = {} + node.parent = parent if label: - parent.metadata["label"] = label + node.metadata["label"] = label if breakdown: - parent.metadata["breakdown"] = breakdown - return parent + node.metadata["breakdown"] = breakdown + if breakdown_labels: + node.metadata["breakdown_labels"] = breakdown_labels + return node def create_mock_scale( diff --git a/tests/test_parameter_labels.py b/tests/test_parameter_labels.py index b055e77..f058441 100644 --- a/tests/test_parameter_labels.py +++ b/tests/test_parameter_labels.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock from policyengine.utils.parameter_labels import ( + _find_breakdown_parent, _generate_bracket_label, _generate_breakdown_label, build_scale_lookup, @@ -407,3 +408,487 @@ def test__given_complex_bracket_path__then_extracts_correct_scale_name( # Then assert result == "Puerto Rico tax rate (bracket 1 rate)" + + +class TestSingleLevelBreakdownParam: + """Tests for single-level breakdown parameters (direct children of breakdown parent).""" + + def test__given_single_level_breakdown_with_enum__then_generates_label_with_enum_value( + self, + ): + # Given: A parameter that is a direct child of a breakdown parent + parent = create_mock_parent_node( + name="gov.tax.exemptions.by_status", + label="Tax exemption by filing status", + breakdown=["filing_status"], + ) + param = create_mock_parameter( + name="gov.tax.exemptions.by_status.MARRIED_FILING_JOINTLY", + parent=parent, + ) + system = create_mock_system( + variables={"filing_status": VARIABLE_WITH_FILING_STATUS_ENUM} + ) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Label uses enum display value + assert result == "Tax exemption by filing status (Married filing jointly)" + + def test__given_single_level_breakdown_without_enum__then_generates_label_with_raw_key( + self, + ): + # Given: A parameter whose breakdown variable has no enum + parent = create_mock_parent_node( + name="gov.benefits.by_category", + label="Benefit by category", + breakdown=["benefit_category"], + ) + param = create_mock_parameter( + name="gov.benefits.by_category.FOOD_ASSISTANCE", + parent=parent, + ) + system = create_mock_system(variables={}) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Falls back to raw key + assert result == "Benefit by category (FOOD_ASSISTANCE)" + + def test__given_single_level_range_breakdown__then_uses_raw_number(self): + # Given: A single-level breakdown with range (no breakdown_labels) + parent = create_mock_parent_node( + name="gov.benefits.by_size", + label="Benefit by household size", + breakdown=["range(1, 8)"], + ) + param = create_mock_parameter( + name="gov.benefits.by_size.4", + parent=parent, + ) + system = create_mock_system() + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Uses raw number without semantic label + assert result == "Benefit by household size (4)" + + def test__given_single_level_range_with_breakdown_label__then_includes_semantic_label( + self, + ): + # Given: A single-level breakdown with range and breakdown_labels + parent = create_mock_parent_node( + name="gov.benefits.by_size", + label="Benefit by household size", + breakdown=["range(1, 8)"], + breakdown_labels=["Household size"], + ) + param = create_mock_parameter( + name="gov.benefits.by_size.4", + parent=parent, + ) + system = create_mock_system() + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Includes semantic label + assert result == "Benefit by household size (Household size 4)" + + +class TestSingleLevelScaleParam: + """Tests confirming scale/bracket params are unaffected by breakdown changes.""" + + def test__given_single_bracket_param__then_generates_bracket_label(self): + # Given: A bracket parameter (not a breakdown) + param = create_mock_parameter(name="gov.tax.income.rates[0].rate") + scale = create_mock_scale( + name="gov.tax.income.rates", + label="Income tax rates", + scale_type="marginal_rate", + ) + system = create_mock_system() + scale_lookup = {"gov.tax.income.rates": scale} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Generates bracket label, not breakdown label + assert result == "Income tax rates (bracket 1 rate)" + + def test__given_bracket_param_with_breakdown_parent__then_bracket_takes_precedence( + self, + ): + # Given: A bracket parameter even if it has a parent with breakdown + # (This shouldn't happen in practice, but tests the precedence) + parent = create_mock_parent_node( + name="gov.tax.rates", + label="Tax rates", + breakdown=["filing_status"], + ) + param = create_mock_parameter( + name="gov.tax.rates[0].rate", + parent=parent, + ) + scale = create_mock_scale( + name="gov.tax.rates", + label="Tax rates", + scale_type="marginal_rate", + ) + system = create_mock_system() + scale_lookup = {"gov.tax.rates": scale} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Bracket notation takes precedence + assert result == "Tax rates (bracket 1 rate)" + + +class TestFindBreakdownParent: + """Tests for the _find_breakdown_parent function.""" + + def test__given_direct_child__then_returns_parent(self): + # Given + parent = create_mock_parent_node( + name="gov.snap.max_allotment", + label="SNAP max allotment", + breakdown=["snap_region", "range(1, 9)"], + ) + param = create_mock_parameter( + name="gov.snap.max_allotment.CONTIGUOUS_US", + parent=parent, + ) + + # When + result = _find_breakdown_parent(param) + + # Then + assert result == parent + + def test__given_nested_child__then_returns_breakdown_ancestor(self): + # Given + breakdown_parent = create_mock_parent_node( + name="gov.snap.max_allotment", + label="SNAP max allotment", + breakdown=["snap_region", "range(1, 9)"], + ) + intermediate = create_mock_parent_node( + name="gov.snap.max_allotment.CONTIGUOUS_US", + parent=breakdown_parent, + ) + param = create_mock_parameter( + name="gov.snap.max_allotment.CONTIGUOUS_US.1", + parent=intermediate, + ) + + # When + result = _find_breakdown_parent(param) + + # Then + assert result == breakdown_parent + + def test__given_no_breakdown_in_ancestry__then_returns_none(self): + # Given + parent = create_mock_parent_node( + name="gov.tax.rate", + label="Tax rate", + ) + param = create_mock_parameter( + name="gov.tax.rate.value", + parent=parent, + ) + + # When + result = _find_breakdown_parent(param) + + # Then + assert result is None + + +class TestNestedBreakdownLabels: + """Tests for nested breakdown label generation.""" + + def test__given_nested_breakdown_with_enum_and_range__then_generates_full_label( + self, + ): + # Given + breakdown_parent = create_mock_parent_node( + name="gov.snap.max_allotment", + label="SNAP max allotment", + breakdown=["snap_region", "range(1, 9)"], + breakdown_labels=["SNAP region", "Household size"], + ) + intermediate = create_mock_parent_node( + name="gov.snap.max_allotment.CONTIGUOUS_US", + parent=breakdown_parent, + ) + param = create_mock_parameter( + name="gov.snap.max_allotment.CONTIGUOUS_US.1", + parent=intermediate, + ) + system = create_mock_system() + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + # Without snap_region enum in system, uses breakdown_label for first dimension too + assert result == "SNAP max allotment (SNAP region CONTIGUOUS_US, Household size 1)" + + def test__given_breakdown_labels_for_range__then_includes_semantic_label( + self, + ): + # Given + parent = create_mock_parent_node( + name="gov.benefits.amount", + label="Benefit amount", + breakdown=["range(1, 5)"], + breakdown_labels=["Number of dependants"], + ) + param = create_mock_parameter( + name="gov.benefits.amount.3", + parent=parent, + ) + system = create_mock_system() + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result == "Benefit amount (Number of dependants 3)" + + def test__given_enum_with_breakdown_label__then_prefers_enum_value(self): + # Given + parent = create_mock_parent_node( + name="gov.exemptions.personal", + label="Personal exemption", + breakdown=["filing_status"], + breakdown_labels=["Filing status"], + ) + param = create_mock_parameter( + name="gov.exemptions.personal.SINGLE", + parent=parent, + ) + system = create_mock_system( + variables={"filing_status": VARIABLE_WITH_FILING_STATUS_ENUM} + ) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + # Enum value "Single" is used, not "Filing status SINGLE" + assert result == "Personal exemption (Single)" + + def test__given_three_level_nesting__then_generates_all_dimensions(self): + # Given + breakdown_parent = create_mock_parent_node( + name="gov.irs.sales_tax", + label="State sales tax", + breakdown=["state_code", "range(1, 7)", "range(1, 20)"], + breakdown_labels=["State", "Income bracket", "Exemption count"], + ) + level1 = create_mock_parent_node( + name="gov.irs.sales_tax.CA", + parent=breakdown_parent, + ) + level2 = create_mock_parent_node( + name="gov.irs.sales_tax.CA.3", + parent=level1, + ) + param = create_mock_parameter( + name="gov.irs.sales_tax.CA.3.5", + parent=level2, + ) + system = create_mock_system( + variables={"state_code": VARIABLE_WITH_STATE_CODE_ENUM} + ) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result == "State sales tax (CA, Income bracket 3, Exemption count 5)" + + def test__given_missing_breakdown_labels__then_uses_raw_values(self): + # Given + breakdown_parent = create_mock_parent_node( + name="gov.snap.max_allotment", + label="SNAP max allotment", + breakdown=["snap_region", "range(1, 9)"], + # No breakdown_labels provided + ) + intermediate = create_mock_parent_node( + name="gov.snap.max_allotment.CONTIGUOUS_US", + parent=breakdown_parent, + ) + param = create_mock_parameter( + name="gov.snap.max_allotment.CONTIGUOUS_US.1", + parent=intermediate, + ) + system = create_mock_system() + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + # Falls back to raw values when no breakdown_labels + assert result == "SNAP max allotment (CONTIGUOUS_US, 1)" + + +class TestMixedNestedBreakdowns: + """Tests for complex mixed breakdowns with enums, ranges, and labels.""" + + def test__given_enum_range_enum_nesting__then_formats_each_correctly(self): + # Given: A complex 3-level breakdown: enum -> range -> enum + # Example: state -> household_size -> filing_status + breakdown_parent = create_mock_parent_node( + name="gov.tax.credits.earned_income", + label="Earned income credit", + breakdown=["state_code", "range(1, 6)", "filing_status"], + breakdown_labels=["State", "Number of children", "Filing status"], + ) + level1 = create_mock_parent_node( + name="gov.tax.credits.earned_income.CA", + parent=breakdown_parent, + ) + level2 = create_mock_parent_node( + name="gov.tax.credits.earned_income.CA.2", + parent=level1, + ) + param = create_mock_parameter( + name="gov.tax.credits.earned_income.CA.2.SINGLE", + parent=level2, + ) + system = create_mock_system( + variables={ + "state_code": VARIABLE_WITH_STATE_CODE_ENUM, + "filing_status": VARIABLE_WITH_FILING_STATUS_ENUM, + } + ) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Enum values use display names, range uses breakdown_label + assert result == "Earned income credit (CA, Number of children 2, Single)" + + def test__given_range_enum_range_nesting__then_formats_each_correctly(self): + # Given: range -> enum -> range nesting + breakdown_parent = create_mock_parent_node( + name="gov.childcare.subsidy", + label="Childcare subsidy", + breakdown=["range(1, 4)", "filing_status", "range(1, 10)"], + breakdown_labels=["Age group", "Filing status", "Household size"], + ) + level1 = create_mock_parent_node( + name="gov.childcare.subsidy.2", + parent=breakdown_parent, + ) + level2 = create_mock_parent_node( + name="gov.childcare.subsidy.2.HEAD_OF_HOUSEHOLD", + parent=level1, + ) + param = create_mock_parameter( + name="gov.childcare.subsidy.2.HEAD_OF_HOUSEHOLD.5", + parent=level2, + ) + system = create_mock_system( + variables={"filing_status": VARIABLE_WITH_FILING_STATUS_ENUM} + ) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Both ranges use breakdown_labels, enum uses display value + assert ( + result + == "Childcare subsidy (Age group 2, Head of household, Household size 5)" + ) + + def test__given_partial_breakdown_labels__then_uses_labels_where_available(self): + # Given: breakdown_labels list shorter than breakdown list + breakdown_parent = create_mock_parent_node( + name="gov.benefits.utility", + label="Utility allowance", + breakdown=["area_code", "range(1, 20)", "housing_type"], + breakdown_labels=["Area", "Household size"], # Missing label for housing_type + ) + level1 = create_mock_parent_node( + name="gov.benefits.utility.AREA_1", + parent=breakdown_parent, + ) + level2 = create_mock_parent_node( + name="gov.benefits.utility.AREA_1.3", + parent=level1, + ) + param = create_mock_parameter( + name="gov.benefits.utility.AREA_1.3.RENTER", + parent=level2, + ) + system = create_mock_system(variables={}) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: Uses breakdown_labels where available, raw value for missing label + assert result == "Utility allowance (Area AREA_1, Household size 3, RENTER)" + + def test__given_four_level_nesting_with_mixed_types__then_generates_all_dimensions( + self, + ): + # Given: A deeply nested 4-level breakdown + breakdown_parent = create_mock_parent_node( + name="gov.irs.deductions.sales_tax", + label="State sales tax deduction", + breakdown=["state_code", "filing_status", "range(1, 7)", "range(1, 20)"], + breakdown_labels=["State", "Filing status", "Exemption count", "Income bracket"], + ) + level1 = create_mock_parent_node( + name="gov.irs.deductions.sales_tax.NY", + parent=breakdown_parent, + ) + level2 = create_mock_parent_node( + name="gov.irs.deductions.sales_tax.NY.JOINT", + parent=level1, + ) + level3 = create_mock_parent_node( + name="gov.irs.deductions.sales_tax.NY.JOINT.4", + parent=level2, + ) + param = create_mock_parameter( + name="gov.irs.deductions.sales_tax.NY.JOINT.4.12", + parent=level3, + ) + system = create_mock_system( + variables={ + "state_code": VARIABLE_WITH_STATE_CODE_ENUM, + "filing_status": VARIABLE_WITH_FILING_STATUS_ENUM, + } + ) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then: All 4 dimensions are formatted correctly + assert ( + result + == "State sales tax deduction (NY, Joint, Exemption count 4, Income bracket 12)" + )