chore: Lint and changelog

Author: anth-volk

changelog_entry.yaml

@@ -0,0 +1,8 @@
+- bump: patch
+  changes:
+    changed:
+    - Disambiguated filepath management in Simulation._set_data()
+    - Refactored Simulation._set_data() to divide functionality into smaller methods
+    - Prevented passage of non-Path URIs to Dataset.from_file() at end of Simulation._set_data() execution
+    added:
+    - Tests for Simulation._set_data()

policyengine/simulation.py

@@ -3,7 +3,12 @@
 import sys
 from pydantic import BaseModel, Field
 from typing import Literal
-from .utils.data.datasets import get_default_dataset, process_gs_path, POLICYENGINE_DATASETS, DATASET_TIME_PERIODS
+from .utils.data.datasets import (
+    get_default_dataset,
+    process_gs_path,
+    POLICYENGINE_DATASETS,
+    DATASET_TIME_PERIODS,
+)
 from policyengine_core.simulations import Simulation as CountrySimulation
 from policyengine_core.simulations import (
     Microsimulation as CountryMicrosimulation,
@@ -31,9 +36,7 @@
 
 CountryType = Literal["uk", "us"]
 ScopeType = Literal["household", "macro"]
-DataType = (
-    str | Dataset | None
-)
+DataType = str | Dataset | None
 TimePeriodType = int
 ReformType = ParametricReform | Type[StructuralReform] | None
 RegionType = Optional[str]
@@ -95,7 +98,7 @@ def __init__(self, **options: SimulationOptions):
         self.options = SimulationOptions(**options)
         self.check_model_version()
         if not isinstance(self.options.data, Dataset):
-          self._set_data(self.options.data)
+            self._set_data(self.options.data)
         self._initialise_simulations()
         self.check_data_version()
         self._add_output_functions()
@@ -139,8 +142,7 @@ def _set_data(self, file_address: str | None = None) -> None:
         # If None is passed, user wants default dataset; get URL, then continue initializing.
         if file_address is None:
             file_address = get_default_dataset(
-                country=self.options.country,
-                region=self.options.region
+                country=self.options.country, region=self.options.region
             )
             print(
                 f"No data provided, using default dataset: {file_address}",
@@ -155,15 +157,11 @@ def _set_data(self, file_address: str | None = None) -> None:
         else:
             # All official PolicyEngine datasets are stored in GCS;
             # load accordingly
-            filename, version = self._set_data_from_gs(
-                file_address
-            )
+            filename, version = self._set_data_from_gs(file_address)
             self.data_version = version
 
-        time_period = self._set_data_time_period(
-            file_address
-        )
-            
+        time_period = self._set_data_time_period(file_address)
+
         self.options.data = Dataset.from_file(
             filename, time_period=time_period
         )
@@ -370,7 +368,7 @@ def check_data_version(self) -> None:
                 raise ValueError(
                     f"Data version {self.data_version} does not match expected version {self.options.data_version}."
                 )
-            
+
     def _set_data_time_period(self, file_address: str) -> Optional[int]:
         """
         Set the time period based on the file address.
@@ -383,9 +381,7 @@ def _set_data_time_period(self, file_address: str) -> Optional[int]:
             # Local file, no time period available
             return None
 
-    def _set_data_from_gs(
-        self, file_address: str
-    ) -> tuple[str, str | None]:
+    def _set_data_from_gs(self, file_address: str) -> tuple[str, str | None]:
         """
         Set the data from a GCS path and return the filename and version.
         """
@@ -403,4 +399,3 @@ def _set_data_from_gs(
         )
 
         return filename, version
-    

policyengine/utils/data/datasets.py

@@ -23,6 +23,7 @@
     ECPS_2024: 2023,
 }
 
+
 def get_default_dataset(
     country: str, region: str, version: Optional[str] = None
 ) -> str:
@@ -38,11 +39,12 @@ def get_default_dataset(
         f"Unable to select a default dataset for country {country} and region {region}."
     )
 
+
 def process_gs_path(path: str) -> Tuple[str, str]:
     """Process a GS path to return bucket and object."""
     if not path.startswith("gs://"):
         raise ValueError(f"Invalid GS path: {path}")
-    
+
     path = path[5:]  # Remove 'gs://'
     bucket, obj = path.split("/", 1)
-    return bucket, obj
+    return bucket, obj

tests/fixtures/simulation.py

@@ -1,7 +1,7 @@
 from policyengine.simulation import SimulationOptions
 from unittest.mock import patch, Mock
 import pytest
-from policyengine.utils.data.datasets import CPS_2023 
+from policyengine.utils.data.datasets import CPS_2023
 
 non_data_uk_sim_options = {
     "country": "uk",
@@ -21,41 +21,43 @@
     "baseline": None,
 }
 
-uk_sim_options_no_data = SimulationOptions.model_validate({
-    **non_data_uk_sim_options,
-    "data": None,
-})
+uk_sim_options_no_data = SimulationOptions.model_validate(
+    {
+        **non_data_uk_sim_options,
+        "data": None,
+    }
+)
 
-us_sim_options_cps_dataset = SimulationOptions.model_validate({
-    **non_data_us_sim_options,
-    "data": CPS_2023
-})
+us_sim_options_cps_dataset = SimulationOptions.model_validate(
+    {**non_data_us_sim_options, "data": CPS_2023}
+)
 
 SAMPLE_DATASET_FILENAME = "sample_value.h5"
 SAMPLE_DATASET_BUCKET_NAME = "policyengine-uk-data-private"
 SAMPLE_DATASET_URI_PREFIX = "gs://"
 SAMPLE_DATASET_FILE_ADDRESS = f"{SAMPLE_DATASET_URI_PREFIX}{SAMPLE_DATASET_BUCKET_NAME}/{SAMPLE_DATASET_FILENAME}"
 
-uk_sim_options_pe_dataset = SimulationOptions.model_validate({
-    **non_data_uk_sim_options,
-    "data": SAMPLE_DATASET_FILE_ADDRESS
-})
+uk_sim_options_pe_dataset = SimulationOptions.model_validate(
+    {**non_data_uk_sim_options, "data": SAMPLE_DATASET_FILE_ADDRESS}
+)
+
 
 @pytest.fixture
 def mock_get_default_dataset():
     with patch(
         "policyengine.simulation.get_default_dataset",
-        return_value=SAMPLE_DATASET_FILE_ADDRESS
+        return_value=SAMPLE_DATASET_FILE_ADDRESS,
     ) as mock_get_default_dataset:
         yield mock_get_default_dataset
 
+
 @pytest.fixture
 def mock_dataset():
     """Simple Dataset mock fixture"""
-    with patch('policyengine.simulation.Dataset') as mock_dataset_class:
+    with patch("policyengine.simulation.Dataset") as mock_dataset_class:
         mock_instance = Mock()
         # Set file_path to mimic Dataset's behavior of clipping URI and bucket name from GCS paths
         mock_instance.from_file = Mock()
         mock_instance.file_path = SAMPLE_DATASET_FILENAME
         mock_dataset_class.from_file.return_value = mock_instance
-        yield mock_instance
+        yield mock_instance

tests/test_simulation.py

@@ -4,52 +4,69 @@
     us_sim_options_cps_dataset,
     mock_get_default_dataset,
     mock_dataset,
-    SAMPLE_DATASET_FILENAME
+    SAMPLE_DATASET_FILENAME,
 )
 import sys
 from copy import deepcopy
 
 from policyengine import Simulation
 
+
 class TestSimulation:
     class TestSetData:
-        def test__given_no_data_option__sets_default_dataset(self, mock_get_default_dataset, mock_dataset):
+        def test__given_no_data_option__sets_default_dataset(
+            self, mock_get_default_dataset, mock_dataset
+        ):
 
             # Don't run entire init script
             sim = object.__new__(Simulation)
             sim.options = deepcopy(uk_sim_options_no_data)
             sim._set_data(uk_sim_options_no_data.data)
 
             assert str(sim.options.data.file_path) == SAMPLE_DATASET_FILENAME
-        def test__given_pe_dataset__sets_data_option_to_dataset(self, mock_dataset):
+
+        def test__given_pe_dataset__sets_data_option_to_dataset(
+            self, mock_dataset
+        ):
 
             sim = object.__new__(Simulation)
             sim.options = deepcopy(uk_sim_options_pe_dataset)
             sim._set_data(uk_sim_options_pe_dataset.data)
 
             assert str(sim.options.data.file_path) == SAMPLE_DATASET_FILENAME
-        def test__given_cps_2023_in_filename__sets_time_period_to_2023(self, mock_dataset):
+
+        def test__given_cps_2023_in_filename__sets_time_period_to_2023(
+            self, mock_dataset
+        ):
             from policyengine import Simulation
 
             sim = object.__new__(Simulation)
             sim.options = deepcopy(us_sim_options_cps_dataset)
             sim._set_data(us_sim_options_cps_dataset.data)
 
             assert mock_dataset.from_file.called_with(
-                us_sim_options_cps_dataset.data,
-                time_period=2023
+                us_sim_options_cps_dataset.data, time_period=2023
             )
+
     class TestSetDataTimePeriod:
         def test__given_dataset_with_time_period__sets_time_period(self):
             from policyengine import Simulation
 
             sim = object.__new__(Simulation)
 
             print("Dataset:", us_sim_options_cps_dataset.data, file=sys.stderr)
-            assert sim._set_data_time_period(us_sim_options_cps_dataset.data) == 2023
+            assert (
+                sim._set_data_time_period(us_sim_options_cps_dataset.data)
+                == 2023
+            )
 
-        def test__given_dataset_without_time_period__does_not_set_time_period(self):
+        def test__given_dataset_without_time_period__does_not_set_time_period(
+            self,
+        ):
             from policyengine import Simulation
 
             sim = object.__new__(Simulation)
-            assert sim._set_data_time_period(uk_sim_options_pe_dataset.data) == None
+            assert (
+                sim._set_data_time_period(uk_sim_options_pe_dataset.data)
+                == None
+            )