Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 120 additions & 14 deletions src/policyengine/utils/parameter_labels.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
18 changes: 12 additions & 6 deletions tests/fixtures/parameter_labels_fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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(
Expand Down
Loading