Forward branches through Social Security period discovery

changelog.d/known-periods-branch-wrapper.fixed.md

@@ -0,0 +1 @@
+Fixed Social Security period discovery to forward explicit branches while preserving compatibility with Core versions that accept no branch argument.

policyengine_us/tests/core/_social_security_state_backend.py

@@ -1,6 +1,6 @@
 """Small storage backend for the Social Security input state machine.
 
-The country is intentionally never imported. Core's installed source bodies
+The country is intentionally never imported. Core's selected source bodies
 implement every write, builder, copy, replay, extraction and dump route below;
 the US module's complete function bodies implement the hooks. Entity and
 variable descriptors, dataset construction, and disk dump bytes are synthetic.
@@ -49,7 +49,6 @@
 from dataclasses import dataclass
 from enum import Enum
 from functools import lru_cache, wraps
-from importlib.metadata import distribution
 from pathlib import Path
 from unittest.mock import patch
 from weakref import WeakKeyDictionary, ref
@@ -179,7 +178,7 @@ def __init__(self, *, cache_blacklist=(), variables_to_drop=()):
     def _setup(self):
         cache_blacklist = self.cache_blacklist
         variables_to_drop = self.variables_to_drop
-        core = Path(distribution("policyengine-core").locate_file("policyengine_core"))
+        core = Path(core_simulation.__file__).resolve().parents[1]
         calculation_frames = core_simulation._calculation_frames
         formula_simulation = core_simulation._formula_simulation
         branch_clone = core_simulation._branch_clone
@@ -389,15 +388,15 @@ def subsample(self, *args):
             def calculate(self, name, period):
                 # The country-free fixture has no tracer or concurrent
                 # formulas. Retain Core's public reader cache behavior while
-                # its installed _calculate body supplies the actual value.
+                # its selected _calculate body supplies the actual value.
                 period = as_period(period)
                 key = (name, period)
                 if key not in self._fast_cache:
                     self._fast_cache[key] = self._calculate(name, period)
                 return self._fast_cache[key]
 
             def _calculate(self, name, period):
-                raise AssertionError("Installed Core calculation body was not loaded")
+                raise AssertionError("Selected Core calculation body was not loaded")
 
             def check_macro_cache(self, *args):
                 return False
@@ -590,12 +589,12 @@ def load(self):
             variable.set_input = core_helpers["set_input_divide_by_period"]
         # Core 3.33.12 cleans overlapping raw seeds after a direct set_input;
         # 3.33.2 does this only in dispatch/divide helpers. Keep the reference
-        # model tied to the installed behavior rather than a version string.
+        # model tied to the selected behavior rather than a version string.
         self.cleans_raw_overlaps = "_drop_calculated_over_own_period" in core_helpers
         simulation_source = core / "simulations/simulation.py"
         # _latest_input_key calls _end_order, which uses Core's calendar
         # table even for yearly sources. Load that literal from the same
-        # installed module rather than maintaining a second calendar table.
+        # selected module rather than maintaining a second calendar table.
         day_offsets = next(
             node.value
             for node in source_tree(simulation_source).body
@@ -640,6 +639,15 @@ def load(self):
             "_input_set_meanwhile",
             "_drop_computed",
         )
+        sim_methods += tuple(
+            name
+            for name in (
+                "_set_fallback_cache_status",
+                "_fallback_is_uncached",
+                "_uncached_fallback",
+            )
+            if hasattr(core_simulation.Simulation, name)
+        )
         for name, method in execute(
             core / "simulations/simulation.py", sim_methods, namespace, "Simulation"
         ).items():

policyengine_us/tests/core/test_social_security_input_state_machine.py

@@ -8,7 +8,7 @@
 reported amount and public formula amount, including the initial scored year
 and a later year whose intermediate base may already have been cached.
 
-The backend executes installed Core source bodies and the production US hooks
+The backend executes selected Core source bodies and the production US hooks
 without importing the country. Run this file directly for a country-free test:
 ``python policyengine_us/tests/core/test_social_security_input_state_machine.py``.
 """

policyengine_us/tests/core/test_social_security_known_periods.py

@@ -0,0 +1,79 @@
+"""Keep reported source periods compatible with both Core holder APIs."""
+
+import inspect
+
+import numpy as np
+import pytest
+from policyengine_core.holders import Holder
+from policyengine_core.periods import period
+
+import test_social_security_real_core as real_core
+from test_social_security_real_core import installed_hooks as installed_hooks
+
+
+@pytest.fixture
+def branch_aware_core():
+    physical = real_core.hooks._PHYSICAL["get_known_periods"]
+    if "branch_name" not in inspect.signature(physical).parameters:
+        pytest.skip("Core master Holder.get_known_periods accepts no branch argument")
+
+
+@pytest.mark.parametrize("public", real_core.COMPONENTS)
+def test_known_periods_without_branch_keeps_reported_sources(public):
+    reported = f"{public}_reported"
+    simulation = real_core.fresh(
+        {public: {"2024": [12_000]}, reported: {"2026": [6_000]}}
+    )
+    holder = simulation.get_holder(reported)
+    assert {str(known) for known in holder.get_known_periods()} == {"2024", "2026"}
+    assert {str(known) for known in holder.get_known_periods(None)} == {"2024", "2026"}
+
+
+@pytest.mark.usefixtures("branch_aware_core")
+@pytest.mark.parametrize("public", real_core.COMPONENTS)
+def test_known_periods_lists_only_sources_readable_by_requested_branch(public):
+    reported = f"{public}_reported"
+    simulation = real_core.fresh()
+    simulation.get_holder(public).set_input(
+        period(2024), np.array([12_000]), branch_name="later"
+    )
+    holder = simulation.get_holder(reported)
+    holder.set_input(period(2026), np.array([6_000]), branch_name="later")
+
+    assert holder.get_array(period(2024), "later")[0] == 12_000
+    assert holder.get_array(period(2024), "default") is None
+    assert {str(known) for known in holder.get_known_periods("later")} == {
+        "2024",
+        "2026",
+    }
+    assert holder.get_known_periods("default") == []
+    assert holder.get_known_periods("unrelated") == []
+
+
+@pytest.mark.usefixtures("branch_aware_core")
+def test_known_periods_forwards_branch_to_fallback_check(monkeypatch):
+    simulation = real_core.fresh()
+    holder = simulation.get_holder(f"{real_core.PUBLIC}_reported")
+    original = real_core.hooks._allow_fallback
+    branches = []
+
+    def allow_fallback(simulation, name, period=None, branch_name=None):
+        if period is None:
+            branches.append(branch_name)
+        return original(simulation, name, period, branch_name)
+
+    monkeypatch.setattr(real_core.hooks, "_allow_fallback", allow_fallback)
+    assert holder.get_known_periods("later") == []
+    assert branches == ["later"]
+
+
+@pytest.mark.parametrize("public", real_core.COMPONENTS)
+def test_root_fallback_calculates_missing_reported_input(public):
+    simulation = real_core.fresh()
+    assert simulation.calculate(f"{public}_reported", 2024)[0] == 0
+
+
+def test_root_fallback_calculates_missing_nonreported_input():
+    simulation = real_core.fresh()
+    simulation.delete_arrays("independent_input")
+    assert simulation.calculate("independent_input", 2024)[0] == 0

policyengine_us/tests/core/test_social_security_real_core.py

@@ -443,6 +443,9 @@ def test_dataset_rebuild_keeps_new_input_provenance_when_old_cache_shares_key():
     assert since is not None
     assert scalar(simulation, "independent_input", "2023") == 0
     old_holder = simulation.get_holder("independent_input")
+    # Core may leave a missing input's default uncached. Store the derived
+    # entry this test needs under the key the dataset's input then reuses.
+    old_holder.put_in_cache(np.array([0.0]), period("2023"), derived=True)
     assert old_holder.is_derived(period("2023"))
     assert sequence_number(simulation, "independent_input", "2023") >= since
 

policyengine_us/tools/social_security_inputs.py

@@ -893,20 +893,27 @@ def get_array(holder, period, branch_name="default"):
     Holder.get_array = get_array
 
     @wraps(_PHYSICAL["get_known_periods"])
-    def get_known_periods(holder):
-        values = _PHYSICAL["get_known_periods"](holder)
+    def get_known_periods(holder, branch_name=None):
+        branch_kwargs = {} if branch_name is None else {"branch_name": branch_name}
+        values = _PHYSICAL["get_known_periods"](holder, **branch_kwargs)
         if holder.variable.name in _REPORTED_COMPONENTS and _enabled(holder.simulation):
-            if _allow_fallback(holder.simulation, holder.variable.name):
+            if _allow_fallback(
+                holder.simulation, holder.variable.name, **branch_kwargs
+            ):
                 values = [
                     *values,
                     *(
                         period
                         for period in _public_periods(
                             holder.simulation,
                             holder.variable.name.removesuffix("_reported"),
+                            **branch_kwargs,
                         )
                         if _effective_value(
-                            holder.simulation, holder.variable.name, period
+                            holder.simulation,
+                            holder.variable.name,
+                            period,
+                            **branch_kwargs,
                         )[0]
                         is not None
                     ),