Merge pull request #209 from PolicyEngine/fix/apply-labels-to-nested-breakdowns

feat: Add support for breakdown_labels and nested breakdowns

Author: anth-volk

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):

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(