Require pure Medicare enrollment input support

changelog.d/pure-medicare-enrollment-input.md

@@ -0,0 +1 @@
+- Require `medicare_enrolled` to be a pure PolicyEngine US input before CPS writes Medicare enrollment values.

policyengine_us_data/utils/policyengine.py

@@ -121,10 +121,15 @@ def assert_locked_policyengine_us_version() -> PolicyEngineUSBuildInfo:
 
 
 @lru_cache(maxsize=1)
-def _policyengine_us_variable_names() -> frozenset[str]:
+def _policyengine_us_variables():
     from policyengine_us import CountryTaxBenefitSystem
 
-    return frozenset(CountryTaxBenefitSystem().variables)
+    return CountryTaxBenefitSystem().variables
+
+
+@lru_cache(maxsize=1)
+def _policyengine_us_variable_names() -> frozenset[str]:
+    return frozenset(_policyengine_us_variables())
 
 
 def has_policyengine_us_variables(*variables: str) -> bool:
@@ -136,8 +141,31 @@ def has_policyengine_us_variables(*variables: str) -> bool:
     return set(variables).issubset(available_variables)
 
 
+def has_policyengine_us_pure_input_variables(*variables: str) -> bool:
+    try:
+        available_variables = _policyengine_us_variables()
+    except Exception:
+        return False
+
+    for variable in variables:
+        variable_definition = available_variables.get(variable)
+        if variable_definition is None:
+            return False
+        if not variable_definition.is_input_variable():
+            return False
+        if getattr(variable_definition, "defined_for", None) is not None:
+            return False
+        if getattr(variable_definition, "adds", None):
+            return False
+        if getattr(variable_definition, "subtracts", None):
+            return False
+        if getattr(variable_definition, "formulas", None):
+            return False
+    return True
+
+
 def supports_medicare_enrollment_input() -> bool:
-    return has_policyengine_us_variables("medicare_enrolled")
+    return has_policyengine_us_pure_input_variables("medicare_enrolled")
 
 
 def supports_modeled_medicare_part_b_inputs() -> bool:

tests/unit/test_medicare_part_b_inputs.py

@@ -6,6 +6,21 @@
 from policyengine_us_data.utils import policyengine as policyengine_utils
 
 
+class _Variable:
+    adds = None
+    defined_for = None
+    formulas = None
+    subtracts = None
+
+    def __init__(self, *, is_input=True, defined_for=None, formulas=None):
+        self._is_input = is_input
+        self.defined_for = defined_for
+        self.formulas = formulas
+
+    def is_input_variable(self):
+        return self._is_input
+
+
 def test_medicare_part_b_clone_imputation_matches_installed_model_support():
     assert ("medicare_part_b_premiums" in set(CPS_ONLY_IMPUTED_VARIABLES)) is (
         not supports_modeled_medicare_part_b_inputs()
@@ -22,6 +37,35 @@ def test_supports_medicare_enrollment_input_allows_partial_support(monkeypatch):
         "has_policyengine_us_variables",
         lambda *variables: variables == ("medicare_enrolled",),
     )
+    monkeypatch.setattr(
+        policyengine_utils,
+        "has_policyengine_us_pure_input_variables",
+        lambda *variables: variables == ("medicare_enrolled",),
+    )
 
     assert policyengine_utils.supports_medicare_enrollment_input() is True
     assert policyengine_utils.supports_modeled_medicare_part_b_inputs() is False
+
+
+def test_pure_input_support_rejects_formula_and_conditional_variables(monkeypatch):
+    monkeypatch.setattr(
+        policyengine_utils,
+        "_policyengine_us_variables",
+        lambda: {
+            "pure_input": _Variable(),
+            "formula_input": _Variable(formulas={"2025": object()}),
+            "conditional_input": _Variable(defined_for="is_eligible"),
+            "formula_variable": _Variable(is_input=False),
+        },
+    )
+
+    assert policyengine_utils.has_policyengine_us_pure_input_variables("pure_input")
+    assert not policyengine_utils.has_policyengine_us_pure_input_variables(
+        "formula_input"
+    )
+    assert not policyengine_utils.has_policyengine_us_pure_input_variables(
+        "conditional_input"
+    )
+    assert not policyengine_utils.has_policyengine_us_pure_input_variables(
+        "formula_variable"
+    )