Use repo-local raw PUF fallback

Author: MaxGhenis

src/microplex_us/data_sources/puf.py

@@ -401,6 +401,31 @@ def download_puf(cache_dir: Path | None = None) -> Path:
     return Path(puf_path), Path(demo_path)
 
 
+def _resolve_policyengine_repo_local_puf_paths(
+    policyengine_us_data_repo: str | Path | None,
+) -> tuple[Path, Path | None] | None:
+    """Resolve raw PUF CSVs from a local policyengine-us-data checkout."""
+
+    if policyengine_us_data_repo is None:
+        return None
+    try:
+        repo_root = resolve_policyengine_us_data_repo_root(policyengine_us_data_repo)
+    except FileNotFoundError:
+        return None
+    candidate_dirs = (
+        repo_root / "policyengine_us_data" / "storage",
+        repo_root / "data" / "raw",
+    )
+    for candidate_dir in candidate_dirs:
+        puf_path = candidate_dir / "puf_2015.csv"
+        demographics_path = candidate_dir / "demographics_2015.csv"
+        if puf_path.exists():
+            return puf_path, (
+                demographics_path if demographics_path.exists() else None
+            )
+    return None
+
+
 def load_puf_raw(puf_path: Path, demographics_path: Path | None = None) -> pd.DataFrame:
     """Load raw PUF data from CSV."""
     print(f"Loading PUF from {puf_path}...")
@@ -1772,8 +1797,16 @@ def load_puf(
     Returns:
         DataFrame with common variable names, ready for stacking with CPS
     """
-    # Download if needed
-    puf_path, demo_path = download_puf(cache_dir)
+    # Prefer a repo-local raw PUF copy when available to avoid remote auth/cache
+    # requirements during rebuild runs.
+    local_repo_paths = _resolve_policyengine_repo_local_puf_paths(
+        policyengine_us_data_repo
+    )
+    if local_repo_paths is not None:
+        puf_path, demo_path = local_repo_paths
+        print(f"Using repo-local PUF from {puf_path}...")
+    else:
+        puf_path, demo_path = download_puf(cache_dir)
 
     # Load raw data
     raw = load_puf_raw(puf_path, demo_path)
@@ -2195,8 +2228,15 @@ def load_frame(self, query: SourceQuery | None = None) -> ObservationFrame:
             )
         )
         if puf_path is None:
-            loader = self.loader or download_puf
-            loaded_puf_path, loaded_demographics_path = loader(self.cache_dir)
+            local_repo_paths = _resolve_policyengine_repo_local_puf_paths(
+                policyengine_us_data_repo
+            )
+            if local_repo_paths is not None:
+                loaded_puf_path, loaded_demographics_path = local_repo_paths
+                print(f"Using repo-local PUF from {loaded_puf_path}...")
+            else:
+                loader = self.loader or download_puf
+                loaded_puf_path, loaded_demographics_path = loader(self.cache_dir)
             puf_path = loaded_puf_path
             if demographics_path is None:
                 demographics_path = loaded_demographics_path

tests/test_puf_source_provider.py

@@ -1058,6 +1058,46 @@ def fail_download(*args, **kwargs):
     assert resolved_demo_path == demographics_path
 
 
+def test_puf_source_provider_prefers_policyengine_repo_local_raw_files(
+    tmp_path, monkeypatch
+):
+    repo_root = tmp_path / "policyengine-us-data"
+    storage_dir = repo_root / "policyengine_us_data" / "storage"
+    storage_dir.mkdir(parents=True)
+    pd.DataFrame(
+        {
+            "RECID": [101],
+            "MARS": [1],
+            "XTOT": [1],
+            "S006": [100.0],
+            "E00200": [50_000.0],
+            "E02400": [0.0],
+            "E01400": [0.0],
+            "AGE_HEAD": [45],
+            "GENDER": [1],
+        }
+    ).to_csv(storage_dir / "puf_2015.csv", index=False)
+    pd.DataFrame({"RECID": [101]}).to_csv(
+        storage_dir / "demographics_2015.csv", index=False
+    )
+
+    def fail_loader(*args, **kwargs):
+        raise AssertionError("remote/cache loader should not run when repo-local PUF exists")
+
+    provider = PUFSourceProvider(
+        target_year=2015,
+        policyengine_us_data_repo=repo_root,
+        loader=fail_loader,
+        social_security_share_model_loader=_mock_social_security_share_model_loader,
+    )
+
+    frame = provider.load_frame(SourceQuery(period=2015))
+    persons = frame.tables[EntityType.PERSON]
+
+    assert len(persons) == 1
+    assert persons["employment_income"].sum() == 50_000.0
+
+
 def test_map_puf_variables_seed_controls_age_imputation():
     puf = pd.DataFrame(
         {