1- """Budget-window annual result extraction and aggregation helpers."""
1+ """Budget-window result validation and aggregation helpers."""
22
33from __future__ import annotations
44
55from decimal import Decimal
66from typing import Any
77
88from src .modal .gateway .models import (
9- BudgetWindowAnnualImpact ,
109 BudgetWindowResult ,
1110 BudgetWindowTotals ,
11+ SingleYearMacroOutput ,
1212)
1313
1414# The UK microsimulation has no state/province fiscal layer, so worker child
2020 "benefit_spending_impact" ,
2121 "budgetary_impact" ,
2222)
23- OPTIONAL_BUDGET_KEYS = ("state_tax_revenue_impact" ,)
23+
24+
25+ def _is_number (value : Any ) -> bool :
26+ return isinstance (value , int | float ) and not isinstance (value , bool )
2427
2528
2629def _as_decimal (value : float | int ) -> Decimal :
@@ -32,46 +35,60 @@ def _as_decimal(value: float | int) -> Decimal:
3235 return Decimal (str (value ))
3336
3437
35- def extract_annual_impact (
38+ def validate_single_year_output (
3639 * ,
3740 simulation_year : str ,
3841 child_result : dict [str , Any ],
39- ) -> BudgetWindowAnnualImpact :
42+ ) -> SingleYearMacroOutput :
43+ """Validate and normalize a child macro result.
44+
45+ UK worker results can omit ``state_tax_revenue_impact`` because there is
46+ no state/province fiscal layer. The canonical output still includes that
47+ field, defaulted to zero, so downstream clients receive one stable shape.
48+ """
49+
50+ if not isinstance (child_result , dict ):
51+ raise ValueError (
52+ "Malformed budget-window child result: expected object for "
53+ f"{ simulation_year }
54+ )
55+
4056 budget = child_result .get ("budget" , {})
4157 if not isinstance (budget , dict ):
4258 raise ValueError ("Malformed budget-window child result: missing budget object" )
4359
4460 missing_keys = [
45- key
46- for key in REQUIRED_BUDGET_KEYS
47- if not isinstance (budget .get (key ), int | float )
61+ key for key in REQUIRED_BUDGET_KEYS if not _is_number (budget .get (key ))
4862 ]
4963 if missing_keys :
5064 missing = ", " .join (f"budget.{ key } for key in missing_keys )
5165 raise ValueError (
5266 f"Malformed budget-window child result: missing numeric { missing }
5367 )
5468
55- tax_revenue_impact = budget ["tax_revenue_impact" ]
56- # UK worker results omit the state fiscal layer entirely; coerce to 0.0
57- # so the parent aggregator can still report federal/state splits with a
58- # uniform shape across countries.
59- state_tax_revenue_impact = budget .get ("state_tax_revenue_impact" )
60- if not isinstance (state_tax_revenue_impact , int | float ):
61- state_tax_revenue_impact = 0.0
62-
63- return BudgetWindowAnnualImpact (
64- year = simulation_year ,
65- taxRevenueImpact = tax_revenue_impact ,
66- federalTaxRevenueImpact = tax_revenue_impact - state_tax_revenue_impact ,
67- stateTaxRevenueImpact = state_tax_revenue_impact ,
68- benefitSpendingImpact = budget ["benefit_spending_impact" ],
69- budgetaryImpact = budget ["budgetary_impact" ],
70- )
69+ normalized = dict (child_result )
70+ normalized_budget = dict (budget )
71+ if "state_tax_revenue_impact" not in normalized_budget :
72+ normalized_budget ["state_tax_revenue_impact" ] = 0.0
73+ elif not _is_number (normalized_budget ["state_tax_revenue_impact" ]):
74+ raise ValueError (
75+ "Malformed budget-window child result: missing numeric "
76+ "budget.state_tax_revenue_impact"
77+ )
78+ normalized ["budget" ] = normalized_budget
79+
80+ try :
81+ return SingleYearMacroOutput .model_validate (normalized )
82+ except Exception as exc :
83+ raise ValueError (
84+ f"Malformed budget-window child result for { simulation_year } { exc }
85+ ) from exc
7186
7287
73- def sum_annual_impacts (
74- annual_impacts : list [BudgetWindowAnnualImpact ],
88+ def sum_single_year_outputs (
89+ * ,
90+ outputs_by_year : dict [str , SingleYearMacroOutput ],
91+ years : list [str ],
7592) -> BudgetWindowTotals :
7693 """Sum per-year impacts using Decimal accumulators.
7794
@@ -93,18 +110,21 @@ def sum_annual_impacts(
93110 "budgetaryImpact" : Decimal (0 ),
94111 }
95112
96- for annual_impact in annual_impacts :
97- totals ["taxRevenueImpact" ] += _as_decimal (annual_impact .taxRevenueImpact )
113+ for year in years :
114+ output = outputs_by_year [year ]
115+ budget = output .model_dump (mode = "json" )["budget" ]
116+ tax_revenue_impact = budget ["tax_revenue_impact" ]
117+ state_tax_revenue_impact = budget .get ("state_tax_revenue_impact" )
118+
119+ totals ["taxRevenueImpact" ] += _as_decimal (tax_revenue_impact )
98120 totals ["federalTaxRevenueImpact" ] += _as_decimal (
99- annual_impact .federalTaxRevenueImpact
100- )
101- totals ["stateTaxRevenueImpact" ] += _as_decimal (
102- annual_impact .stateTaxRevenueImpact
121+ tax_revenue_impact - state_tax_revenue_impact
103122 )
123+ totals ["stateTaxRevenueImpact" ] += _as_decimal (state_tax_revenue_impact )
104124 totals ["benefitSpendingImpact" ] += _as_decimal (
105- annual_impact . benefitSpendingImpact
125+ budget [ "benefit_spending_impact" ]
106126 )
107- totals ["budgetaryImpact" ] += _as_decimal (annual_impact . budgetaryImpact )
127+ totals ["budgetaryImpact" ] += _as_decimal (budget [ "budgetary_impact" ] )
108128
109129 return BudgetWindowTotals (** {key : float (value ) for key , value in totals .items ()})
110130
@@ -113,12 +133,25 @@ def build_budget_window_result(
113133 * ,
114134 start_year : str ,
115135 window_size : int ,
116- annual_impacts : list [ BudgetWindowAnnualImpact ],
136+ outputs_by_year : dict [ str , SingleYearMacroOutput ],
117137) -> BudgetWindowResult :
138+ years = [str (int (start_year ) + offset ) for offset in range (window_size )]
139+ missing_years = [year for year in years if year not in outputs_by_year ]
140+ if missing_years :
141+ raise ValueError (
142+ "Cannot build budget-window result: missing outputs for "
143+ + ", " .join (missing_years )
144+ )
145+
146+ ordered_outputs = {year : outputs_by_year [year ] for year in years }
118147 return BudgetWindowResult (
119148 startYear = start_year ,
120149 endYear = str (int (start_year ) + window_size - 1 ),
121150 windowSize = window_size ,
122- annualImpacts = annual_impacts ,
123- totals = sum_annual_impacts (annual_impacts ),
151+ years = years ,
152+ outputsByYear = ordered_outputs ,
153+ totals = sum_single_year_outputs (
154+ outputs_by_year = ordered_outputs ,
155+ years = years ,
156+ ),
124157 )
0 commit comments