Merge branch 'codex/fix-private-default-dataset' of https://github.com/PolicyEngine/policyengine-uk into codex/merge-1549

Author: MaxGhenis

policyengine_uk/simulation.py

@@ -1,4 +1,5 @@
 # Standard library imports
+import os
 from typing import Dict, Optional, Union, Type, List
 
 # Third-party imports
@@ -27,7 +28,7 @@
 from policyengine_uk.utils.dependencies import get_variable_dependencies
 from policyengine_uk.reforms import create_structural_reforms_from_parameters
 
-from .tax_benefit_system import CountryTaxBenefitSystem
+from .tax_benefit_system import CountryTaxBenefitSystem, DEFAULT_DATASET_ENV_VAR
 
 from microdf import MicroDataFrame
 
@@ -37,6 +38,16 @@
 _url_dataset_cache: dict = {}
 
 
+def get_default_dataset_url() -> str:
+    dataset_url = os.environ.get(DEFAULT_DATASET_ENV_VAR)
+    if dataset_url:
+        return dataset_url
+    raise ValueError(
+        "Simulation() requires an explicit dataset when no situation is provided. "
+        f"Pass dataset=..., or set {DEFAULT_DATASET_ENV_VAR} to opt into a default dataset."
+    )
+
+
 def _pre_encode_enum_columns(
     dataset: UKMultiYearDataset, tbs: "CountryTaxBenefitSystem"
 ) -> None:
@@ -143,9 +154,7 @@ def __init__(
         elif isinstance(dataset, UKMultiYearDataset):
             self.build_from_multi_year_dataset(dataset)
         elif dataset is None:
-            self.build_from_url(
-                "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2023_24.h5"
-            )
+            self.build_from_url(get_default_dataset_url())
         else:
             raise ValueError(f"Unsupported dataset type: {dataset.__class__}")
 

policyengine_uk/tax_benefit_system.py

@@ -37,7 +37,7 @@
 
 # Module constants
 COUNTRY_DIR = Path(__file__).parent
-ENHANCED_FRS = "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2023_24.h5"
+DEFAULT_DATASET_ENV_VAR = "POLICYENGINE_UK_DEFAULT_DATASET"
 
 # Cache for fully-processed parameter tree, so convert_to_fiscal_year_parameters
 # (22,538 param.update() calls) only runs once per process.

policyengine_uk/tests/conftest.py

@@ -0,0 +1,28 @@
+import os
+
+import pytest
+
+DEFAULT_TEST_DATASET_URL = (
+    "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3"
+)
+
+if os.environ.get("HUGGING_FACE_TOKEN") and not os.environ.get(
+    "POLICYENGINE_UK_DEFAULT_DATASET"
+):
+    os.environ["POLICYENGINE_UK_DEFAULT_DATASET"] = DEFAULT_TEST_DATASET_URL
+
+
+def pytest_collection_modifyitems(config, items):
+    has_default_dataset = bool(os.environ.get("POLICYENGINE_UK_DEFAULT_DATASET"))
+    if has_default_dataset:
+        return
+
+    skip_microsimulation = pytest.mark.skip(
+        reason=(
+            "Requires POLICYENGINE_UK_DEFAULT_DATASET or HUGGING_FACE_TOKEN "
+            "for microsimulation dataset access"
+        )
+    )
+    for item in items:
+        if "microsimulation" in item.keywords:
+            item.add_marker(skip_microsimulation)

policyengine_uk/tests/test_behavioral_responses.py

@@ -16,8 +16,11 @@
 from policyengine_uk import Microsimulation
 from policyengine_uk.model_api import Scenario
 
-# Check if HF token is available for data-dependent tests
-HF_TOKEN_AVAILABLE = bool(os.environ.get("HUGGING_FACE_TOKEN"))
+# Check if a default dataset is available for data-dependent tests
+HF_TOKEN_AVAILABLE = bool(
+    os.environ.get("HUGGING_FACE_TOKEN")
+    or os.environ.get("POLICYENGINE_UK_DEFAULT_DATASET")
+)
 requires_hf_data = pytest.mark.skipif(
     not HF_TOKEN_AVAILABLE,
     reason="Requires HUGGING_FACE_TOKEN for private data access",

policyengine_uk/tests/test_deterministic_variables.py

@@ -155,9 +155,15 @@ def test_household_owns_tv_can_be_set_false(self):
 
 
 class TestDefaultDatasetUrl:
-    """Test that the default dataset URL points at the private HF repo."""
+    """Test explicit handling of the default dataset URL."""
 
-    def test_simulation_defaults_to_private_hf_repo(self, monkeypatch):
+    def test_simulation_requires_explicit_default_dataset(self, monkeypatch):
+        monkeypatch.delenv("POLICYENGINE_UK_DEFAULT_DATASET", raising=False)
+
+        with pytest.raises(ValueError, match="requires an explicit dataset"):
+            Simulation()
+
+    def test_simulation_uses_opt_in_default_dataset(self, monkeypatch):
         captured = {}
 
         class _StopDefaultDatasetLoad(Exception):
@@ -168,12 +174,16 @@ def fake_build_from_url(self, url):
             raise _StopDefaultDatasetLoad
 
         monkeypatch.setattr(Simulation, "build_from_url", fake_build_from_url)
+        monkeypatch.setenv(
+            "POLICYENGINE_UK_DEFAULT_DATASET",
+            "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5",
+        )
 
         with pytest.raises(_StopDefaultDatasetLoad):
             Simulation()
 
         assert captured["url"] == (
-            "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2023_24.h5"
+            "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5"
         )