From 18cb84d4fd0712408b826a7d05bd295f4b517f28 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Sat, 24 Oct 2026 01:04:06 +0100 Subject: [PATCH 1/2] Validate zero-based integer enum positions --- changelog.d/fix-integer-enum-bounds.fixed.md | 1 + policyengine_core/enums/enum.py | 53 ++++++- policyengine_core/variables/variable.py | 2 + tests/core/enums/test_integer_enum_bounds.py | 150 +++++++++++++++++++ tests/core/test_enum_input_bounds.py | 125 ++++++++++++++++ 5 files changed, 325 insertions(+), 6 deletions(-) create mode 100644 changelog.d/fix-integer-enum-bounds.fixed.md create mode 100644 tests/core/enums/test_integer_enum_bounds.py create mode 100644 tests/core/test_enum_input_bounds.py diff --git a/changelog.d/fix-integer-enum-bounds.fixed.md b/changelog.d/fix-integer-enum-bounds.fixed.md new file mode 100644 index 000000000..f5a1a93e7 --- /dev/null +++ b/changelog.d/fix-integer-enum-bounds.fixed.md @@ -0,0 +1 @@ +- Validate zero-based integer enum positions before narrowing, including encoded arrays and situation inputs, and report the enum in invalid-position errors. diff --git a/policyengine_core/enums/enum.py b/policyengine_core/enums/enum.py index 8d91e691e..44575f4f7 100644 --- a/policyengine_core/enums/enum.py +++ b/policyengine_core/enums/enum.py @@ -39,30 +39,70 @@ def _get_sorted_lookup_arrays(cls) -> Tuple[np.ndarray, np.ndarray]: ) return sorted_names_arr, sorted_indices + @classmethod + def _validate_integer_indices(cls, values) -> None: + """Validate zero-based positions before a scalar or array is narrowed.""" + if isinstance(values, (int, np.integer)): + if isinstance(values, (bool, np.bool_)): + return + if 0 <= values < len(cls): + return + invalid_values = [int(values)] + else: + raw = ( + values.view(np.ndarray) + if isinstance(values, EnumArray) + else np.asarray(values) + ) + if raw.dtype.kind not in {"i", "u"}: + return + invalid = (raw < 0) | (raw >= len(cls)) + if not np.any(invalid): + return + invalid_values = np.unique(raw[invalid]).tolist() + + shown = invalid_values[:5] + remainder = len(invalid_values) - len(shown) + suffix = f" (and {remainder} more)" if remainder else "" + raise ValueError( + f"Invalid position(s) {shown}{suffix} for enum {cls.__name__}. " + f"Integer positions are zero-based and must be in [0, {len(cls)})." + ) + @classmethod def encode(cls, array: Union[EnumArray, np.ndarray]) -> EnumArray: """ - Encode an array of enum items or string identifiers into an EnumArray. + Encode enum items, names, or zero-based integer positions into an EnumArray. Args: array: The input array to encode. Can be an EnumArray, a NumPy array - of enum items, or a NumPy array of string identifiers. + of enum items, string identifiers, or integer positions. Integer + positions follow declaration order, from zero through one less + than the number of members; they are not survey source codes. Returns: An EnumArray containing the encoded values. + Raises: + ValueError: An integer position or string identifier is invalid. + Examples: + >>> class MyEnum(Enum): + ... ITEM_1 = "First" + ... ITEM_2 = "Second" + ... ITEM_3 = "Third" >>> string_array = np.array(["ITEM_1", "ITEM_2", "ITEM_3"]) >>> encoded_array = MyEnum.encode(string_array) - >>> encoded_array - EnumArray([1, 2, 3], dtype=int8) + >>> encoded_array.view(np.ndarray).tolist() + [0, 1, 2] >>> item_array = np.array([MyEnum.ITEM_1, MyEnum.ITEM_2, MyEnum.ITEM_3]) >>> encoded_array = MyEnum.encode(item_array) - >>> encoded_array - EnumArray([1, 2, 3], dtype=int8) + >>> encoded_array.view(np.ndarray).tolist() + [0, 1, 2] """ if isinstance(array, EnumArray): + cls._validate_integer_indices(array) return array # Handle Enum item arrays by extracting indices directly @@ -104,6 +144,7 @@ def encode(cls, array: Union[EnumArray, np.ndarray]) -> EnumArray: ) elif array.dtype.kind in {"i", "u"}: # Integer array - already indices + cls._validate_integer_indices(array) indices = array else: raise ValueError(f"Unsupported array dtype: {array.dtype}") diff --git a/policyengine_core/variables/variable.py b/policyengine_core/variables/variable.py index de9415c42..2c82e49fb 100644 --- a/policyengine_core/variables/variable.py +++ b/policyengine_core/variables/variable.py @@ -763,6 +763,8 @@ def clone(self): return clone def check_set_value(self, value): + if self.value_type == Enum and isinstance(value, (int, numpy.integer)): + self.possible_values._validate_integer_indices(value) if self.value_type == Enum and isinstance(value, str): try: value = self.possible_values[value].index diff --git a/tests/core/enums/test_integer_enum_bounds.py b/tests/core/enums/test_integer_enum_bounds.py new file mode 100644 index 000000000..1472d7ca5 --- /dev/null +++ b/tests/core/enums/test_integer_enum_bounds.py @@ -0,0 +1,150 @@ +"""Integer enum positions are zero-based and validated before narrowing.""" + +import numpy as np +import pytest + +from policyengine_core.enums import Enum, EnumArray +from policyengine_core.enums.config import ENUM_ARRAY_DTYPE + + +class PositionChoice(Enum): + FIRST = "First" + SECOND = "Second" + LAST = "Last" + + +INTEGER_DTYPES = [ + pytest.param(np.int8, id="int8"), + pytest.param(np.int16, id="int16"), + pytest.param(np.int32, id="int32"), + pytest.param(np.int64, id="int64"), + pytest.param(np.uint8, id="uint8"), + pytest.param(np.uint16, id="uint16"), + pytest.param(np.uint32, id="uint32"), + pytest.param(np.uint64, id="uint64"), +] + +INVALID_ARRAYS = [ + pytest.param(np.array([-1], dtype=np.int8), "-1", id="negative-int8"), + pytest.param( + np.array([np.iinfo(np.int64).min], dtype=np.int64), + str(np.iinfo(np.int64).min), + id="signed-minimum", + ), + pytest.param(np.array([3], dtype=np.int8), "3", id="upper-bound-int8"), + pytest.param(np.array([3], dtype=np.uint8), "3", id="upper-bound-uint8"), + pytest.param(np.array([65536], dtype=np.int64), "65536", id="wide-signed"), + pytest.param(np.array([65535], dtype=np.uint16), "65535", id="wide-unsigned"), + pytest.param( + np.array([np.iinfo(np.uint64).max], dtype=np.uint64), + str(np.iinfo(np.uint64).max), + id="unsigned-maximum", + ), + pytest.param(np.array([0, -1, 2, 4], dtype=np.int64), "4", id="mixed-bounds"), +] + + +@pytest.mark.parametrize("dtype", INTEGER_DTYPES) +def test_integer_positions_keep_zero_based_order(dtype): + values = np.array([0, 2, 1, 0], dtype=dtype) + original = values.copy() + encoded = PositionChoice.encode(values) + assert isinstance(encoded, EnumArray) + assert encoded.possible_values is PositionChoice + assert encoded.dtype == np.dtype(ENUM_ARRAY_DTYPE) + np.testing.assert_array_equal(encoded.view(np.ndarray), [0, 2, 1, 0]) + np.testing.assert_array_equal( + encoded.decode_to_str(), ["FIRST", "LAST", "SECOND", "FIRST"] + ) + np.testing.assert_array_equal(values, original) + + +@pytest.mark.parametrize("values,bad_value", INVALID_ARRAYS) +def test_invalid_integer_positions_raise_contextual_value_error(values, bad_value): + original = values.copy() + with pytest.raises(ValueError) as error: + PositionChoice.encode(values) + assert "PositionChoice" in str(error.value) + assert bad_value in str(error.value) + np.testing.assert_array_equal(values, original) + + +@pytest.mark.parametrize("dtype", [np.int64, np.uint64], ids=["signed", "unsigned"]) +def test_empty_typed_integer_array_is_still_supported(dtype): + encoded = PositionChoice.encode(np.array([], dtype=dtype)) + assert isinstance(encoded, EnumArray) + assert encoded.possible_values is PositionChoice + assert encoded.shape == (0,) + assert encoded.dtype == np.dtype(ENUM_ARRAY_DTYPE) + + +@pytest.mark.parametrize("dtype", [np.int16, np.uint64], ids=["stored", "wide"]) +def test_valid_same_enum_array_retains_identity(dtype): + values = EnumArray(np.array([0, 2], dtype=dtype), PositionChoice) + assert PositionChoice.encode(values) is values + assert values.possible_values is PositionChoice + assert values.dtype == np.dtype(dtype) + np.testing.assert_array_equal(values.view(np.ndarray), [0, 2]) + + +@pytest.mark.parametrize( + "values", + [ + pytest.param(np.array([-1], dtype=np.int16), id="negative"), + pytest.param(np.array([3], dtype=np.int16), id="upper-bound"), + pytest.param(np.array([65536], dtype=np.int64), id="wide-signed"), + pytest.param( + np.array([np.iinfo(np.uint64).max], dtype=np.uint64), + id="unsigned-maximum", + ), + ], +) +def test_enum_array_cannot_bypass_integer_bounds(values): + encoded = EnumArray(values.copy(), PositionChoice) + with pytest.raises(ValueError, match="PositionChoice"): + PositionChoice.encode(encoded) + assert encoded.possible_values is PositionChoice + np.testing.assert_array_equal(encoded.view(np.ndarray), values) + + +def test_valid_foreign_enum_array_keeps_existing_identity_without_relabelling(): + class OtherChoice(Enum): + ONE = "One" + TWO = "Two" + + values = EnumArray(np.array([0, 1], dtype=np.int16), OtherChoice) + assert PositionChoice.encode(values) is values + assert values.possible_values is OtherChoice + + +def test_valid_unbound_enum_array_is_not_silently_relabelled(): + values = EnumArray(np.array([0, 2], dtype=np.int16)) + assert PositionChoice.encode(values) is values + assert values.possible_values is None + + +@pytest.mark.parametrize( + "values", + [ + pytest.param(np.array(["FIRST", "LAST"]), id="names"), + pytest.param(np.array([b"FIRST", b"LAST"]), id="utf8-bytes"), + pytest.param( + np.array([PositionChoice.FIRST, PositionChoice.LAST]), id="enum-items" + ), + ], +) +def test_existing_name_and_item_routes_keep_their_positions(values): + encoded = PositionChoice.encode(values) + np.testing.assert_array_equal(encoded.view(np.ndarray), [0, 2]) + + +@pytest.mark.parametrize( + "values", + [ + pytest.param(np.array([0.5]), id="float"), + pytest.param(np.array([True]), id="boolean"), + ], +) +def test_integer_bounds_do_not_add_new_dtype_coercions(values): + with pytest.raises(ValueError, match="Unsupported array dtype"): + PositionChoice.encode(values) diff --git a/tests/core/test_enum_input_bounds.py b/tests/core/test_enum_input_bounds.py new file mode 100644 index 000000000..78ec6c3a7 --- /dev/null +++ b/tests/core/test_enum_input_bounds.py @@ -0,0 +1,125 @@ +"""Public enum-input boundaries reject bad positions before dtype narrowing.""" + +import numpy as np +import pytest + +from policyengine_core.entities import build_entity +from policyengine_core.enums import Enum +from policyengine_core.errors import SituationParsingError +from policyengine_core.periods import YEAR +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.variables import Variable + + +class PositionChoice(Enum): + FIRST = "First" + SECOND = "Second" + LAST = "Last" + + +Person = build_entity("person", "persons", "Person", is_person=True) + + +class reported_position(Variable): + value_type = Enum + possible_values = PositionChoice + default_value = PositionChoice.FIRST + entity = Person + definition_period = YEAR + label = "Reported synthetic position" + + +INVALID_SCALARS = [ + pytest.param(-1, id="negative"), + pytest.param(3, id="upper-bound"), + pytest.param(np.int64(65536), id="wide-signed"), + pytest.param(np.uint64(65536), id="wide-unsigned"), + pytest.param(np.uint64(np.iinfo(np.uint64).max), id="unsigned-maximum"), +] + + +@pytest.fixture +def position_system(): + system = TaxBenefitSystem([Person]) + system.add_variable(reported_position) + return system + + +@pytest.mark.parametrize("value", INVALID_SCALARS) +def test_variable_checks_integer_bounds_before_narrowing(position_system, value): + variable = position_system.get_variable("reported_position") + with pytest.raises(ValueError, match="PositionChoice"): + variable.check_set_value(value) + + +@pytest.mark.parametrize( + "value,expected", + [ + pytest.param("FIRST", 0, id="name"), + pytest.param(0, 0, id="first-position"), + pytest.param(np.int64(2), 2, id="last-position"), + ], +) +def test_variable_preserves_valid_names_and_positions(position_system, value, expected): + assert ( + position_system.get_variable("reported_position").check_set_value(value) + == expected + ) + + +@pytest.mark.parametrize("value", INVALID_SCALARS) +def test_set_input_rejects_bad_positions_without_replacing_valid_input( + position_system, value +): + simulation = SimulationBuilder().build_from_dict( + position_system, {"persons": {"p": {}}} + ) + simulation.set_input("reported_position", 2026, np.array([1], dtype=np.int64)) + with pytest.raises(ValueError, match="PositionChoice"): + simulation.set_input("reported_position", 2026, np.array([value])) + np.testing.assert_array_equal( + np.asarray(simulation.calculate("reported_position", 2026)), [1] + ) + + +@pytest.mark.parametrize("dtype", [np.int64, np.uint64], ids=["signed", "unsigned"]) +@pytest.mark.parametrize("value", [0, 1, 2], ids=["first", "middle", "last"]) +def test_set_input_accepts_valid_zero_based_positions(position_system, dtype, value): + simulation = SimulationBuilder().build_from_dict( + position_system, {"persons": {"p": {}}} + ) + simulation.set_input("reported_position", 2026, np.array([value], dtype=dtype)) + np.testing.assert_array_equal( + np.asarray(simulation.calculate("reported_position", 2026)), [value] + ) + + +@pytest.mark.parametrize("value", INVALID_SCALARS) +def test_situation_rejects_bad_positions_before_builder_narrowing( + position_system, value +): + situation = {"persons": {"p": {"reported_position": {"2026": value}}}} + with pytest.raises(SituationParsingError) as error: + SimulationBuilder().build_from_dict(position_system, situation) + assert "reported_position" in str(error.value) + assert "PositionChoice" in str(error.value) + + +@pytest.mark.parametrize( + "value,expected", + [ + pytest.param("FIRST", 0, id="name"), + pytest.param(0, 0, id="first-position"), + pytest.param(2, 2, id="last-position"), + pytest.param(np.int64(2), 2, id="numpy-position"), + ], +) +def test_situation_keeps_valid_names_and_positions(position_system, value, expected): + simulation = SimulationBuilder().build_from_dict( + position_system, + {"persons": {"p": {"reported_position": {"2026": value}}}}, + ) + np.testing.assert_array_equal( + np.asarray(simulation.calculate("reported_position", 2026)), [expected] + ) From 9242c26a76a01582306084009b9a338d90a9a63f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Sat, 24 Oct 2026 01:25:17 +0100 Subject: [PATCH 2/2] Validate encoded enum positions against their owner --- policyengine_core/enums/enum.py | 6 +++- tests/core/enums/test_integer_enum_bounds.py | 29 ++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/policyengine_core/enums/enum.py b/policyengine_core/enums/enum.py index 44575f4f7..656673568 100644 --- a/policyengine_core/enums/enum.py +++ b/policyengine_core/enums/enum.py @@ -82,6 +82,8 @@ def encode(cls, array: Union[EnumArray, np.ndarray]) -> EnumArray: Returns: An EnumArray containing the encoded values. + An existing EnumArray keeps its identity and attached enum; positions + are validated against that enum, or the receiving enum when unbound. Raises: ValueError: An integer position or string identifier is invalid. @@ -102,7 +104,9 @@ def encode(cls, array: Union[EnumArray, np.ndarray]) -> EnumArray: [0, 1, 2] """ if isinstance(array, EnumArray): - cls._validate_integer_indices(array) + possible_values = getattr(array, "possible_values", None) + owner = cls if possible_values is None else possible_values + owner._validate_integer_indices(array) return array # Handle Enum item arrays by extracting indices directly diff --git a/tests/core/enums/test_integer_enum_bounds.py b/tests/core/enums/test_integer_enum_bounds.py index 1472d7ca5..ab75d62cd 100644 --- a/tests/core/enums/test_integer_enum_bounds.py +++ b/tests/core/enums/test_integer_enum_bounds.py @@ -148,3 +148,32 @@ def test_existing_name_and_item_routes_keep_their_positions(values): def test_integer_bounds_do_not_add_new_dtype_coercions(values): with pytest.raises(ValueError, match="Unsupported array dtype"): PositionChoice.encode(values) + + +def test_foreign_bound_array_valid_for_its_owner_retains_identity(): + class LongChoice(Enum): + FIRST = "First" + SECOND = "Second" + THIRD = "Third" + FOURTH = "Fourth" + + values = EnumArray(np.array([3], dtype=np.int16), LongChoice) + assert PositionChoice.encode(values) is values + assert values.possible_values is LongChoice + assert values.dtype == np.dtype(np.int16) + np.testing.assert_array_equal(values.view(np.ndarray), [3]) + np.testing.assert_array_equal(values.decode_to_str(), ["FOURTH"]) + + +def test_foreign_bound_array_invalid_for_its_owner_is_rejected(): + class ShortChoice(Enum): + FIRST = "First" + SECOND = "Second" + + values = EnumArray(np.array([2], dtype=np.int16), ShortChoice) + with pytest.raises(ValueError, match="ShortChoice") as error: + PositionChoice.encode(values) + assert "2" in str(error.value) + assert values.possible_values is ShortChoice + assert values.dtype == np.dtype(np.int16) + np.testing.assert_array_equal(values.view(np.ndarray), [2])