Refactor structural mortgage input compatibility

Author: MaxGhenis

policyengine_us_data/datasets/cps/extended_cps.py

@@ -10,17 +10,23 @@
 from policyengine_us_data.datasets.puf import PUF, PUF_2024
 from policyengine_us_data.storage import STORAGE_FOLDER
 from policyengine_us_data.utils.mortgage_interest import (
+    STRUCTURAL_MORTGAGE_VARIABLES,
     convert_mortgage_interest_to_structural_inputs,
     impute_tax_unit_mortgage_balance_hints,
-    supports_structural_mortgage_inputs,
 )
+from policyengine_us_data.utils.policyengine import has_policyengine_us_variables
 from policyengine_us_data.utils.retirement_limits import (
     get_retirement_limits,
     get_se_pension_limits,
 )
 
 logger = logging.getLogger(__name__)
 
+
+def _supports_structural_mortgage_inputs() -> bool:
+    return has_policyengine_us_variables(*STRUCTURAL_MORTGAGE_VARIABLES)
+
+
 # CPS-only variables that should be QRF-imputed for the PUF clone half
 # instead of naively duplicated from the CPS donor. These are
 # income-correlated variables that exist only in the CPS; demographics,
@@ -450,14 +456,15 @@ def generate(self):
         )
 
         new_data = self._rename_imputed_to_inputs(new_data)
-        new_data = impute_tax_unit_mortgage_balance_hints(
-            new_data,
-            self.time_period,
-        )
-        new_data = convert_mortgage_interest_to_structural_inputs(
-            new_data,
-            self.time_period,
-        )
+        if _supports_structural_mortgage_inputs():
+            new_data = impute_tax_unit_mortgage_balance_hints(
+                new_data,
+                self.time_period,
+            )
+            new_data = convert_mortgage_interest_to_structural_inputs(
+                new_data,
+                self.time_period,
+            )
         new_data = self._drop_formula_variables(new_data)
         self.save_dataset(new_data)
 
@@ -492,7 +499,7 @@ def _rename_imputed_to_inputs(cls, data):
     @classmethod
     def _keep_formula_vars(cls):
         keep = set(cls._KEEP_FORMULA_VARS)
-        if not supports_structural_mortgage_inputs():
+        if not _supports_structural_mortgage_inputs():
             keep.add("interest_deduction")
         return keep
 

policyengine_us_data/datasets/puf/puf.py

@@ -13,8 +13,10 @@
     disaggregate_aggregate_records,
 )
 from policyengine_us_data.utils.mortgage_interest import (
+    STRUCTURAL_MORTGAGE_VARIABLES,
     convert_mortgage_interest_to_structural_inputs,
 )
+from policyengine_us_data.utils.policyengine import has_policyengine_us_variables
 from policyengine_us_data.utils.uprating import (
     create_policyengine_uprating_factors_table,
 )
@@ -650,10 +652,11 @@ def generate(self):
             variable: {self.time_period: values}
             for variable, values in self.holder.items()
         }
-        holder_tp = convert_mortgage_interest_to_structural_inputs(
-            holder_tp,
-            self.time_period,
-        )
+        if has_policyengine_us_variables(*STRUCTURAL_MORTGAGE_VARIABLES):
+            holder_tp = convert_mortgage_interest_to_structural_inputs(
+                holder_tp,
+                self.time_period,
+            )
         self.holder = {
             variable: values[self.time_period] for variable, values in holder_tp.items()
         }