From bcb16001a92b92f0929192f812362bfef524e4a7 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 6 May 2025 16:17:06 -0400 Subject: [PATCH 1/5] fix: Convert Infinity and -Infinity to np.inf and -np.inf --- policyengine/utils/reforms.py | 40 ++++++++++-- tests/utils/test_reforms.py | 115 ++++++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 4 deletions(-) create mode 100644 tests/utils/test_reforms.py diff --git a/policyengine/utils/reforms.py b/policyengine/utils/reforms.py index 6d681f57..3a67f379 100644 --- a/policyengine/utils/reforms.py +++ b/policyengine/utils/reforms.py @@ -1,13 +1,45 @@ import re -from pydantic import RootModel, ValidationError, Field, model_validator -from typing import Dict, TYPE_CHECKING -from annotated_types import Ge, Le +from pydantic import ( + RootModel, + ValidationError, + Field, + model_validator, + field_validator, +) +from typing import Dict, Self, Any, TYPE_CHECKING from typing_extensions import Annotated from typing import Callable from policyengine_core.simulations import Simulation +class ParameterChangeDict(RootModel): + """A dict of changes to a parameter, with custom date string as keys + and various possible value types.""" + + root: Dict[str, Any] + + @model_validator(mode="after") + def check_keys(self) -> Self: + for key in self.root.keys(): + # Check if key is YYYY-MM-DD.YYYY-MM-DD + if not re.match(r"^\d{4}-\d{2}-\d{2}\.\d{4}-\d{2}-\d{2}$", key): + raise ValueError(f"Invalid date format in key: {key}") + return self + + # Convert "Infinity" to "np.inf" and "-Infinity" to "-np.inf" + @field_validator("root", mode="after") + @classmethod + def convert_infinity(cls, value: Dict[str, Any]) -> Dict[str, Any]: + for key, val in value.items(): + if isinstance(val, str): + if val == "Infinity": + value[key] = float("inf") + elif val == "-Infinity": + value[key] = float("-inf") + return value + + class ParametricReform(RootModel): """A reform that just changes parameter values.""" - root: Dict[str, Dict | float | bool] + root: Dict[str, ParameterChangeDict] diff --git a/tests/utils/test_reforms.py b/tests/utils/test_reforms.py new file mode 100644 index 00000000..c9de022a --- /dev/null +++ b/tests/utils/test_reforms.py @@ -0,0 +1,115 @@ +from pydantic import ValidationError +import pytest +import numpy as np + +from policyengine.utils.reforms import ParameterChangeDict, ParametricReform + + +class TestParameterChangeDict: + def test_schema__given_float_inputs__returns_valid_dict(self): + input_data = { + "2023-01-01.2023-12-31": 0.1, + "2024-01-01.2024-12-31": 0.2, + } + + result = ParameterChangeDict(root=input_data) + + assert isinstance(result, ParameterChangeDict) + assert result.root == input_data + + def test_schema__given_string_inputs__returns_valid_dict(self): + input_data = { + "2023-01-01.2023-12-31": "0.1", + "2024-01-01.2024-12-31": "0.2", + } + + result = ParameterChangeDict(root=input_data) + + assert isinstance(result, ParameterChangeDict) + assert result.root == input_data + + def test_schema__given_infinity_string__returns_valid_dict(self): + input_data = { + "2023-01-01.2023-12-31": "Infinity", + "2024-01-01.2024-12-31": "-Infinity", + } + + result = ParameterChangeDict(root=input_data) + + assert isinstance(result, ParameterChangeDict) + assert result.root == { + "2023-01-01.2023-12-31": np.inf, + "2024-01-01.2024-12-31": -np.inf, + } + + def test_schema__given_invalid_date_format__raises_validation_error(self): + input_data = {"2023-01-01.2023-12-31": 0.1, "invalid_date_format": 0.2} + + with pytest.raises( + ValidationError, match="Invalid date format in key" + ): + ParameterChangeDict(root=input_data) + + def test_schema__given_invalid_key_type__raises_validation_error(self): + input_data = {123: 0.1, "2024-01-01.2024-12-31": 0.2} + + with pytest.raises( + ValidationError, match="validation error for ParameterChangeDict" + ): + ParameterChangeDict(root=input_data) + + +class TestParametricReform: + def test_schema__given_valid_dict__returns_valid_reform(self): + input_data = { + "parameter1": { + "2023-01-01.2023-12-31": 0.1, + "2024-01-01.2024-12-31": 0.2, + }, + "parameter2": { + "2023-01-01.2023-12-31": 0.3, + "2024-01-01.2024-12-31": 0.4, + }, + } + + expected_output_data = { + "parameter1": ParameterChangeDict( + root={ + "2023-01-01.2023-12-31": 0.1, + "2024-01-01.2024-12-31": 0.2, + } + ), + "parameter2": ParameterChangeDict( + root={ + "2023-01-01.2023-12-31": 0.3, + "2024-01-01.2024-12-31": 0.4, + } + ), + } + + result = ParametricReform(root=input_data) + + assert isinstance(result, ParametricReform) + assert result.root == expected_output_data + + def test_schema__given_invalid_key_type__raises_validation_error(self): + input_data = { + 123: {"2023-01-01.2023-12-31": 0.1, "2024-01-01.2024-12-31": 0.2}, + "valid_parameter": { + "2023-01-01.2023-12-31": 0.3, + "2024-01-01.2024-12-31": 0.4, + }, + } + + with pytest.raises( + ValidationError, match=r"validation errors? for ParametricReform" + ): + ParametricReform(root=input_data) + + def test_schema__given_dateless_structure__raises_validation_error(self): + input_data = {"parameter1": 0.1, "parameter2": 0.2} + + with pytest.raises( + ValidationError, match=r"validation errors? for ParametricReform" + ): + ParametricReform(root=input_data) From 7a3eb9bc7bbf5f7b8ad2e4bd26a1d2cef32b716a Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 7 May 2025 18:03:12 -0400 Subject: [PATCH 2/5] fix: Redo ParametricReform schema --- changelog_entry.yaml | 5 + policyengine/utils/reforms.py | 86 ++++++++++---- tests/utils/test_reforms.py | 211 +++++++++++++++++++++++++++++++--- 3 files changed, 261 insertions(+), 41 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29b..5822f7a9 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,5 @@ +- bump: patch + changes: + changed: + - Refactored ParametricReform schema into clearer subschemas. + - Added conversion of Infinity and -Infinity to np.inf and -np.inf. \ No newline at end of file diff --git a/policyengine/utils/reforms.py b/policyengine/utils/reforms.py index 3a67f379..8b461079 100644 --- a/policyengine/utils/reforms.py +++ b/policyengine/utils/reforms.py @@ -1,45 +1,81 @@ -import re from pydantic import ( RootModel, - ValidationError, Field, - model_validator, field_validator, ) -from typing import Dict, Self, Any, TYPE_CHECKING +from typing import Dict, Any from typing_extensions import Annotated -from typing import Callable from policyengine_core.simulations import Simulation -class ParameterChangeDict(RootModel): - """A dict of changes to a parameter, with custom date string as keys - and various possible value types.""" +class ParameterChangeValue(RootModel): + """A value for a parameter change, which can be any primitive type or 'Infinity'/'-Infinity'""" - root: Dict[str, Any] + # To prevent validation errors, allow all types except containers + # via field validator + root: Any - @model_validator(mode="after") - def check_keys(self) -> Self: - for key in self.root.keys(): - # Check if key is YYYY-MM-DD.YYYY-MM-DD - if not re.match(r"^\d{4}-\d{2}-\d{2}\.\d{4}-\d{2}-\d{2}$", key): - raise ValueError(f"Invalid date format in key: {key}") - return self + @field_validator("root", mode="after") + @classmethod + def check_type(cls, value: Any) -> Any: + # Check if the value is not a container type + if isinstance(value, (dict, list, set, tuple)): + raise ValueError( + "ParameterChangeValue must not be a container type (dict, list, set, tuple)" + ) + return value # Convert "Infinity" to "np.inf" and "-Infinity" to "-np.inf" @field_validator("root", mode="after") @classmethod - def convert_infinity(cls, value: Dict[str, Any]) -> Dict[str, Any]: - for key, val in value.items(): - if isinstance(val, str): - if val == "Infinity": - value[key] = float("inf") - elif val == "-Infinity": - value[key] = float("-inf") + def convert_infinity(cls, value: Any) -> Any: + if isinstance(value, str): + if value == "Infinity": + value = float("inf") + elif value == "-Infinity": + value = float("-inf") return value +class ParameterChangePeriod(RootModel): + """A period for a parameter change, which can be a single year or a date range""" + + root: Annotated[ + str, + Field( + pattern=r"^\d{4}$|^\d{4}-\d{2}-\d{2}\.\d{4}-\d{2}-\d{2}$", + description="A single year (YYYY) or a date range (YYYY-MM-DD.YYYY-MM-DD)", + ), + ] + + def __hash__(self): + return hash(self.root) + + def __eq__(self, other): + if isinstance(other, ParameterChangePeriod): + return self.root == other.root + return False + + +class ParameterChangeDict(RootModel): + """ + A dict of changes to a parameter, with custom date string as keys + and various possible value types. + + Keys can be formatted one of two ways: + 1. A single year (e.g., "YYYY") + 2. A date range (e.g., "YYYY-MM-DD.YYYY-MM-DD") + """ + + root: Dict[ParameterChangePeriod, ParameterChangeValue] + + class ParametricReform(RootModel): - """A reform that just changes parameter values.""" + """ + A reform that just changes parameter values. + + This is a dict that equates a parameter name to either a single value or a dict of changes. + + """ - root: Dict[str, ParameterChangeDict] + root: Dict[str, ParameterChangeValue | ParameterChangeDict] diff --git a/tests/utils/test_reforms.py b/tests/utils/test_reforms.py index c9de022a..c71ba20a 100644 --- a/tests/utils/test_reforms.py +++ b/tests/utils/test_reforms.py @@ -2,7 +2,12 @@ import pytest import numpy as np -from policyengine.utils.reforms import ParameterChangeDict, ParametricReform +from policyengine.utils.reforms import ( + ParameterChangeDict, + ParametricReform, + ParameterChangeValue, + ParameterChangePeriod, +) class TestParameterChangeDict: @@ -12,10 +17,19 @@ def test_schema__given_float_inputs__returns_valid_dict(self): "2024-01-01.2024-12-31": 0.2, } + expected_output_data = { + ParameterChangePeriod( + root="2023-01-01.2023-12-31" + ): ParameterChangeValue(root=0.1), + ParameterChangePeriod( + root="2024-01-01.2024-12-31" + ): ParameterChangeValue(root=0.2), + } + result = ParameterChangeDict(root=input_data) assert isinstance(result, ParameterChangeDict) - assert result.root == input_data + assert result.root == expected_output_data def test_schema__given_string_inputs__returns_valid_dict(self): input_data = { @@ -23,10 +37,19 @@ def test_schema__given_string_inputs__returns_valid_dict(self): "2024-01-01.2024-12-31": "0.2", } + expected_output_data = { + ParameterChangePeriod( + root="2023-01-01.2023-12-31" + ): ParameterChangeValue(root="0.1"), + ParameterChangePeriod( + root="2024-01-01.2024-12-31" + ): ParameterChangeValue(root="0.2"), + } + result = ParameterChangeDict(root=input_data) assert isinstance(result, ParameterChangeDict) - assert result.root == input_data + assert result.root == expected_output_data def test_schema__given_infinity_string__returns_valid_dict(self): input_data = { @@ -34,19 +57,42 @@ def test_schema__given_infinity_string__returns_valid_dict(self): "2024-01-01.2024-12-31": "-Infinity", } + expected_output_data = { + ParameterChangePeriod( + root="2023-01-01.2023-12-31" + ): ParameterChangeValue(root=np.inf), + ParameterChangePeriod( + root="2024-01-01.2024-12-31" + ): ParameterChangeValue(root=-np.inf), + } + result = ParameterChangeDict(root=input_data) assert isinstance(result, ParameterChangeDict) - assert result.root == { - "2023-01-01.2023-12-31": np.inf, - "2024-01-01.2024-12-31": -np.inf, + assert result.root == expected_output_data + + def test_schema__given_yearly_input__returns_valid_dict(self): + input_data = { + "2023": 0.1, + "2024": 0.2, + } + + expected_output_data = { + ParameterChangePeriod(root="2023"): ParameterChangeValue(root=0.1), + ParameterChangePeriod(root="2024"): ParameterChangeValue(root=0.2), } + result = ParameterChangeDict(root=input_data) + + assert isinstance(result, ParameterChangeDict) + assert result.root == expected_output_data + def test_schema__given_invalid_date_format__raises_validation_error(self): input_data = {"2023-01-01.2023-12-31": 0.1, "invalid_date_format": 0.2} with pytest.raises( - ValidationError, match="Invalid date format in key" + ValidationError, + match=r"validation errors? for ParameterChangeDict", ): ParameterChangeDict(root=input_data) @@ -59,8 +105,93 @@ def test_schema__given_invalid_key_type__raises_validation_error(self): ParameterChangeDict(root=input_data) +class TestParameterChangePeriod: + def test_schema__given_valid_year__returns_valid_period(self): + input_data = "2023" + + result = ParameterChangePeriod(root=input_data) + + assert isinstance(result, ParameterChangePeriod) + assert result.root == input_data + + def test_schema__given_valid_date_range__returns_valid_period(self): + input_data = "2023-01-01.2023-12-31" + + result = ParameterChangePeriod(root=input_data) + + assert isinstance(result, ParameterChangePeriod) + assert result.root == input_data + + def test_schema__given_invalid_date_format__raises_validation_error(self): + input_data = "2023.01.01-2024.12.31" + + with pytest.raises( + ValidationError, + match=r"validation errors? for ParameterChangePeriod", + ): + ParameterChangePeriod(root=input_data) + + +class TestParameterChangeValue: + def test_schema__given_float_input__returns_valid_value(self): + input_data = 0.1 + + result = ParameterChangeValue(root=input_data) + + assert isinstance(result, ParameterChangeValue) + assert result.root == input_data + + def test_schema__given_string_input__returns_valid_value(self): + input_data = "0.1" + + result = ParameterChangeValue(root=input_data) + + assert isinstance(result, ParameterChangeValue) + assert result.root == input_data + + def test_schema__given_bool_input__returns_valid_value(self): + input_data = True + + result = ParameterChangeValue(root=input_data) + + assert isinstance(result, ParameterChangeValue) + assert result.root == input_data + + def test_schema__given_infinity_string__returns_valid_value(self): + input_data = "Infinity" + + result = ParameterChangeValue(root=input_data) + + assert isinstance(result, ParameterChangeValue) + assert result.root == float("inf") + + def test_schema__given_negative_infinity_string__returns_valid_value(self): + input_data = "-Infinity" + + result = ParameterChangeValue(root=input_data) + + assert isinstance(result, ParameterChangeValue) + assert result.root == float("-inf") + + def test_schema__given_invalid_type__raises_validation_error(self): + input_data = [0.1, 0.2] + + with pytest.raises( + ValidationError, match="validation error for ParameterChangeValue" + ): + ParameterChangeValue(root=input_data) + + def test_schema__given_dict_input__raises_validation_error(self): + input_data = {"key": "value"} + + with pytest.raises( + ValidationError, match="validation error for ParameterChangeValue" + ): + ParameterChangeValue(root=input_data) + + class TestParametricReform: - def test_schema__given_valid_dict__returns_valid_reform(self): + def test_schema__given_full_date_dict__returns_valid_reform(self): input_data = { "parameter1": { "2023-01-01.2023-12-31": 0.1, @@ -92,6 +223,62 @@ def test_schema__given_valid_dict__returns_valid_reform(self): assert isinstance(result, ParametricReform) assert result.root == expected_output_data + def test_schema__given_yearly_dict__returns_valid_reform(self): + input_data = { + "parameter1": {"2023": 0.1, "2024": 0.2}, + "parameter2": {"2023": 0.3, "2024": 0.4}, + } + + expected_output_data = { + "parameter1": ParameterChangeDict(root={"2023": 0.1, "2024": 0.2}), + "parameter2": ParameterChangeDict(root={"2023": 0.3, "2024": 0.4}), + } + + result = ParametricReform(root=input_data) + + assert isinstance(result, ParametricReform) + assert result.root == expected_output_data + + def test_schema__given_single_value_dict__returns_valid_reform(self): + input_data = { + "parameter1": 0.1, + "parameter2": 0.2, + } + + expected_output_data = { + "parameter1": ParameterChangeValue(root=0.1), + "parameter2": ParameterChangeValue(root=0.2), + } + + result = ParametricReform(root=input_data) + + assert isinstance(result, ParametricReform) + assert result.root == expected_output_data + + def test_schema__given_mixed_dict__returns_valid_reform(self): + input_data = { + "parameter1": { + "2023-01-01.2023-12-31": 0.1, + "2024-01-01.2024-12-31": 0.2, + }, + "parameter2": 0.3, + } + + expected_output_data = { + "parameter1": ParameterChangeDict( + root={ + "2023-01-01.2023-12-31": 0.1, + "2024-01-01.2024-12-31": 0.2, + } + ), + "parameter2": ParameterChangeValue(root=0.3), + } + + result = ParametricReform(root=input_data) + + assert isinstance(result, ParametricReform) + assert result.root == expected_output_data + def test_schema__given_invalid_key_type__raises_validation_error(self): input_data = { 123: {"2023-01-01.2023-12-31": 0.1, "2024-01-01.2024-12-31": 0.2}, @@ -105,11 +292,3 @@ def test_schema__given_invalid_key_type__raises_validation_error(self): ValidationError, match=r"validation errors? for ParametricReform" ): ParametricReform(root=input_data) - - def test_schema__given_dateless_structure__raises_validation_error(self): - input_data = {"parameter1": 0.1, "parameter2": 0.2} - - with pytest.raises( - ValidationError, match=r"validation errors? for ParametricReform" - ): - ParametricReform(root=input_data) From a38aa42462326b2c6072763692f1c57a6c845f91 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 7 May 2025 18:18:01 -0400 Subject: [PATCH 3/5] fix: Remove ParameterChangePeriod type --- policyengine/simulation.py | 1 + policyengine/utils/reforms.py | 43 +++++++++++------------- tests/utils/test_reforms.py | 63 +++++++++++------------------------ 3 files changed, 39 insertions(+), 68 deletions(-) diff --git a/policyengine/simulation.py b/policyengine/simulation.py index 61f00acb..0c9ad622 100644 --- a/policyengine/simulation.py +++ b/policyengine/simulation.py @@ -195,6 +195,7 @@ def _initialise_simulation( }[country][macro] if isinstance(reform, ParametricReform): + print("Dumping reform model") reform = reform.model_dump() simulation: CountrySimulation = _simulation_type( diff --git a/policyengine/utils/reforms.py b/policyengine/utils/reforms.py index 8b461079..0070246a 100644 --- a/policyengine/utils/reforms.py +++ b/policyengine/utils/reforms.py @@ -1,11 +1,8 @@ from pydantic import ( RootModel, - Field, field_validator, ) from typing import Dict, Any -from typing_extensions import Annotated -from policyengine_core.simulations import Simulation class ParameterChangeValue(RootModel): @@ -37,26 +34,6 @@ def convert_infinity(cls, value: Any) -> Any: return value -class ParameterChangePeriod(RootModel): - """A period for a parameter change, which can be a single year or a date range""" - - root: Annotated[ - str, - Field( - pattern=r"^\d{4}$|^\d{4}-\d{2}-\d{2}\.\d{4}-\d{2}-\d{2}$", - description="A single year (YYYY) or a date range (YYYY-MM-DD.YYYY-MM-DD)", - ), - ] - - def __hash__(self): - return hash(self.root) - - def __eq__(self, other): - if isinstance(other, ParameterChangePeriod): - return self.root == other.root - return False - - class ParameterChangeDict(RootModel): """ A dict of changes to a parameter, with custom date string as keys @@ -67,7 +44,25 @@ class ParameterChangeDict(RootModel): 2. A date range (e.g., "YYYY-MM-DD.YYYY-MM-DD") """ - root: Dict[ParameterChangePeriod, ParameterChangeValue] + root: Dict[str, ParameterChangeValue] + + @field_validator("root", mode="after") + @classmethod + def validate_dates( + cls, value: Dict[str, ParameterChangeValue] + ) -> Dict[str, ParameterChangeValue]: + + year_keys_re = r"^\d{4}$" + date_range_keys_re = r"^\d{4}-\d{2}-\d{2}\.\d{4}-\d{2}-\d{2}$" + + for key in value.keys(): + if not year_keys_re.match(key) and not date_range_keys_re.match( + key + ): + raise ValueError( + f"Key '{key}' must be a single year (YYYY) or a date range (YYYY-MM-DD.YYYY-MM-DD)" + ) + return value class ParametricReform(RootModel): diff --git a/tests/utils/test_reforms.py b/tests/utils/test_reforms.py index c71ba20a..f7a30c6c 100644 --- a/tests/utils/test_reforms.py +++ b/tests/utils/test_reforms.py @@ -18,12 +18,8 @@ def test_schema__given_float_inputs__returns_valid_dict(self): } expected_output_data = { - ParameterChangePeriod( - root="2023-01-01.2023-12-31" - ): ParameterChangeValue(root=0.1), - ParameterChangePeriod( - root="2024-01-01.2024-12-31" - ): ParameterChangeValue(root=0.2), + "2023-01-01.2023-12-31": ParameterChangeValue(root=0.1), + "2024-01-01.2024-12-31": ParameterChangeValue(root=0.2), } result = ParameterChangeDict(root=input_data) @@ -38,12 +34,8 @@ def test_schema__given_string_inputs__returns_valid_dict(self): } expected_output_data = { - ParameterChangePeriod( - root="2023-01-01.2023-12-31" - ): ParameterChangeValue(root="0.1"), - ParameterChangePeriod( - root="2024-01-01.2024-12-31" - ): ParameterChangeValue(root="0.2"), + "2023-01-01.2023-12-31": ParameterChangeValue(root="0.1"), + "2024-01-01.2024-12-31": ParameterChangeValue(root="0.2"), } result = ParameterChangeDict(root=input_data) @@ -58,12 +50,8 @@ def test_schema__given_infinity_string__returns_valid_dict(self): } expected_output_data = { - ParameterChangePeriod( - root="2023-01-01.2023-12-31" - ): ParameterChangeValue(root=np.inf), - ParameterChangePeriod( - root="2024-01-01.2024-12-31" - ): ParameterChangeValue(root=-np.inf), + "2023-01-01.2023-12-31": ParameterChangeValue(root=np.inf), + "2024-01-01.2024-12-31": ParameterChangeValue(root=-np.inf), } result = ParameterChangeDict(root=input_data) @@ -78,8 +66,8 @@ def test_schema__given_yearly_input__returns_valid_dict(self): } expected_output_data = { - ParameterChangePeriod(root="2023"): ParameterChangeValue(root=0.1), - ParameterChangePeriod(root="2024"): ParameterChangeValue(root=0.2), + "2023": ParameterChangeValue(root=0.1), + "2024": ParameterChangeValue(root=0.2), } result = ParameterChangeDict(root=input_data) @@ -96,7 +84,7 @@ def test_schema__given_invalid_date_format__raises_validation_error(self): ): ParameterChangeDict(root=input_data) - def test_schema__given_invalid_key_type__raises_validation_error(self): + def test_schema__given_non_date_key_type__raises_validation_error(self): input_data = {123: 0.1, "2024-01-01.2024-12-31": 0.2} with pytest.raises( @@ -104,32 +92,19 @@ def test_schema__given_invalid_key_type__raises_validation_error(self): ): ParameterChangeDict(root=input_data) - -class TestParameterChangePeriod: - def test_schema__given_valid_year__returns_valid_period(self): - input_data = "2023" - - result = ParameterChangePeriod(root=input_data) - - assert isinstance(result, ParameterChangePeriod) - assert result.root == input_data - - def test_schema__given_valid_date_range__returns_valid_period(self): - input_data = "2023-01-01.2023-12-31" - - result = ParameterChangePeriod(root=input_data) - - assert isinstance(result, ParameterChangePeriod) - assert result.root == input_data - - def test_schema__given_invalid_date_format__raises_validation_error(self): - input_data = "2023.01.01-2024.12.31" + def test_schema__given_incorrect_date_key_type__raises_validation_error( + self, + ): + input_data = { + "2023-01-01.2023-12-31": 0.1, + "2024-01-01.2024-12-31": 0.2, + "2024.01.01-2025.12.31": 0.3, + } with pytest.raises( - ValidationError, - match=r"validation errors? for ParameterChangePeriod", + ValidationError, match="validation error for ParameterChangeDict" ): - ParameterChangePeriod(root=input_data) + ParameterChangeDict(root=input_data) class TestParameterChangeValue: From a15145dd9bea2d23abf09d17baeef3fcc9558088 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 7 May 2025 18:23:02 -0400 Subject: [PATCH 4/5] fix: Fix regex --- policyengine/utils/reforms.py | 5 +++-- tests/utils/test_reforms.py | 1 - 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/policyengine/utils/reforms.py b/policyengine/utils/reforms.py index 0070246a..6dfc804b 100644 --- a/policyengine/utils/reforms.py +++ b/policyengine/utils/reforms.py @@ -1,3 +1,4 @@ +import re from pydantic import ( RootModel, field_validator, @@ -56,8 +57,8 @@ def validate_dates( date_range_keys_re = r"^\d{4}-\d{2}-\d{2}\.\d{4}-\d{2}-\d{2}$" for key in value.keys(): - if not year_keys_re.match(key) and not date_range_keys_re.match( - key + if not re.match(year_keys_re, key) and not re.match( + date_range_keys_re, key ): raise ValueError( f"Key '{key}' must be a single year (YYYY) or a date range (YYYY-MM-DD.YYYY-MM-DD)" diff --git a/tests/utils/test_reforms.py b/tests/utils/test_reforms.py index f7a30c6c..cf43630d 100644 --- a/tests/utils/test_reforms.py +++ b/tests/utils/test_reforms.py @@ -6,7 +6,6 @@ ParameterChangeDict, ParametricReform, ParameterChangeValue, - ParameterChangePeriod, ) From 5ac78ee5ae870c7f7eaf78e14e639d0b109e8dfa Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 9 May 2025 20:10:35 -0400 Subject: [PATCH 5/5] fix: Remove unneeded print statement --- policyengine/simulation.py | 1 - 1 file changed, 1 deletion(-) diff --git a/policyengine/simulation.py b/policyengine/simulation.py index 0c9ad622..61f00acb 100644 --- a/policyengine/simulation.py +++ b/policyengine/simulation.py @@ -195,7 +195,6 @@ def _initialise_simulation( }[country][macro] if isinstance(reform, ParametricReform): - print("Dumping reform model") reform = reform.model_dump() simulation: CountrySimulation = _simulation_type(