diff --git a/src/power_grid_model_ds/_core/model/arrays/base/array.py b/src/power_grid_model_ds/_core/model/arrays/base/array.py index 9c309828..8d328848 100644 --- a/src/power_grid_model_ds/_core/model/arrays/base/array.py +++ b/src/power_grid_model_ds/_core/model/arrays/base/array.py @@ -37,6 +37,44 @@ Self = TypeVar("Self", bound="FancyArray") +def _resolve_str_dtype(name: str, dtype: Any, str_lengths: dict[str, int]) -> Any: + """Replace np.str_ with a fixed-length unicode dtype, leaving other dtypes untouched.""" + if dtype is np.str_: + return np.dtype(f"U{str_lengths.get(name, _DEFAULT_STR_LENGTH)}") + return dtype + + +def _parse_annotation_pre_25(name: str, type_def: Any, type_args: tuple, str_lengths: dict[str, int]) -> tuple: + """Parse an NDArray annotation into a numpy dtype tuple for NumPy < 2.5.""" + # Expected type_args for NDArray[]: (tuple[typing.Any, ...], numpy.dtype[numpy.int32]) + if len(type_args) == 2 and get_origin(type_args[1]) is np.dtype: # noqa: PLR2004 + dtype = get_args(type_args[1])[0] + return (name, _resolve_str_dtype(name, dtype, str_lengths)) + # Expected type_args for NDArray3[]: + # (numpy.ndarray[tuple[typing.Any, ...], numpy.dtype[numpy.float64]], typing.Literal[3]) + if len(type_args) == 2 and get_origin(type_args[1]) is Literal: # noqa: PLR2004 + try: + dtype = get_args(get_args(type_args[0])[1])[0] + size = get_args(type_args[1])[0] + except IndexError as error: + raise ValueError(f"dtype {type_def} not understood or supported") from error + return (name, _resolve_str_dtype(name, dtype, str_lengths), size) + raise ValueError(f"dtype {type_def} not understood or supported") + + +def _parse_annotation_post_25(name: str, type_def: Any, type_args: tuple, str_lengths: dict[str, int]) -> tuple: + """Parse an NDArray annotation into a numpy dtype tuple for NumPy >= 2.5.""" + # Expected type_args for NDArray: (numpy.int32,) + if len(type_args) == 1: + dtype = type_args[0] + return (name, _resolve_str_dtype(name, dtype, str_lengths)) + # Expected type_args for NDArray3: (NDArray[numpy.float64], typing.Literal[3]) + if len(type_args) == 2: # noqa: PLR2004 + dtype = get_args(type_args[0])[0] + return (name, _resolve_str_dtype(name, dtype, str_lengths), get_args(type_args[1])[0]) + raise ValueError(f"dtype {type_def} not understood or supported") + + class FancyArray(ABC): # noqa: B024 """Base class for all arrays. @@ -102,45 +140,29 @@ def get_id_columns(cls) -> set[str]: @lru_cache def get_dtype(cls): # noqa: python:S3776 annotations = get_public_annotations(cls) + + if not annotations.keys(): + raise ArrayDefinitionError(f"Array '{cls.__name__}' has no defined Columns") + + if reserved := set(annotations.keys()) & _RESERVED_COLUMN_NAMES: + raise ArrayDefinitionError( + f"Columns of '{cls.__name__}' cannot be reserved names: {reserved} " + f"(reserved names are: {_RESERVED_COLUMN_NAMES})" + ) + str_lengths = combine_attribute_from_parent_classes(cls, "_str_lengths", dict) - dtypes = {} # Numpy 2.5 changed the typing interface, so we need to treat these differently - is_before_numpy_25 = version.parse(np.__version__) < version.parse("2.5.0") - - for name, type_def in annotations.items(): - type_args = get_args(type_def) - - # Expected type_args pre-2.5 for NDArray[]: (tuple[typing.Any, ...], numpy.dtype[numpy.int32]) - if is_before_numpy_25 and len(type_args) == 2 and get_origin(type_args[1]) is np.dtype: # noqa: PLR2004 - dtypes[name] = get_args(type_args[1])[0] - # Expected type_args pre-2.5 for NDArray3[]: - # (numpy.ndarray[tuple[typing.Any, ...], numpy.dtype[numpy.float64]], typing.Literal[3]) - elif is_before_numpy_25 and len(type_args) == 2 and get_origin(type_args[1]) is Literal: # noqa: PLR2004 - dtypes[name] = (get_args(get_args(type_args[0])[1])[0], get_args(type_args[1])[0]) - # Expected type_args post-2.5 for NDArray: (numpy.int32,) - elif len(type_args) == 1: # pragma: no cover - dtypes[name] = type_args[0] - # Expected type_args post-2.5 for NDArray3: (NDArray[numpy.float64], typing.Literal[3]) - elif len(type_args) == 2: # noqa: PLR2004 # pragma: no cover - dtypes[name] = (get_args(type_args[0])[0], get_args(type_args[1])[0]) - else: - raise ValueError(f"dtype {type_def} not understood or supported") - - if not dtypes: - raise ArrayDefinitionError("Array has no defined Columns") - if reserved := set(dtypes.keys()) & _RESERVED_COLUMN_NAMES: - raise ArrayDefinitionError(f"Columns cannot be reserved names: {reserved}") - - dtype_list = [] - for name, dtype in dtypes.items(): - if dtype is np.str_: - string_length = str_lengths.get(name, _DEFAULT_STR_LENGTH) - dtype_list.append((name, np.dtype(f"U{string_length}"))) - elif dtype is tuple: - dtype_list.append((name, *dtype)) - else: - dtype_list.append((name, dtype)) + parse_annotation = ( + _parse_annotation_pre_25 + if version.parse(np.__version__) < version.parse("2.5.0") + else _parse_annotation_post_25 + ) + + dtype_list = [ + parse_annotation(name, type_def, get_args(type_def), str_lengths) for name, type_def in annotations.items() + ] + return np.dtype(dtype_list) def __repr__(self) -> str: diff --git a/tests/unit/model/arrays/test_get_dtype.py b/tests/unit/model/arrays/test_get_dtype.py new file mode 100644 index 00000000..c7f9fefa --- /dev/null +++ b/tests/unit/model/arrays/test_get_dtype.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Contributors to the Power Grid Model project +# +# SPDX-License-Identifier: MPL-2.0 + +from typing import Any, Literal + +import numpy as np +import pytest + +from power_grid_model_ds._core.model.arrays.base import array as array_module +from power_grid_model_ds._core.model.arrays.base.array import ( + _parse_annotation_post_25, + _parse_annotation_pre_25, +) + + +class TestParseAnnotationPre25: + def test_scalar(self): + type_args = (tuple[Any, ...], np.dtype[np.int64]) + assert _parse_annotation_pre_25("value", None, type_args, {}) == ("value", np.int64) + + def test_ndarray3(self): + type_args = (np.ndarray[tuple[Any, ...], np.dtype[np.float64]], Literal[3]) + assert _parse_annotation_pre_25("value", None, type_args, {}) == ("value", np.float64, 3) + + def test_str_length_folding(self): + type_args = (tuple[Any, ...], np.dtype[np.str_]) + assert _parse_annotation_pre_25("name", None, type_args, {"name": 100}) == ("name", np.dtype("U100")) + + def test_str_default_length(self): + type_args = (tuple[Any, ...], np.dtype[np.str_]) + expected = ("name", np.dtype(f"U{array_module._DEFAULT_STR_LENGTH}")) + assert _parse_annotation_pre_25("name", None, type_args, {}) == expected + + def test_unsupported_shape_raises(self): + with pytest.raises(ValueError, match="not understood or supported"): + _parse_annotation_pre_25("value", "bad", (), {}) + + def test_malformed_ndarray3_raises(self): + # A Literal-tagged annotation whose inner element lacks the expected nested structure must + # raise the clear ValueError rather than an opaque IndexError. + type_args = (int, Literal[3]) + with pytest.raises(ValueError, match="not understood or supported"): + _parse_annotation_pre_25("value", "bad", type_args, {}) + + +class TestParseAnnotationPost25: + def test_scalar(self): + assert _parse_annotation_post_25("value", None, (np.int64,), {}) == ("value", np.int64) + + def test_ndarray3(self): + type_args = (np.dtype[np.float64], Literal[3]) + assert _parse_annotation_post_25("value", None, type_args, {}) == ("value", np.float64, 3) + + def test_str_length_folding(self): + assert _parse_annotation_post_25("name", None, (np.str_,), {"name": 100}) == ("name", np.dtype("U100")) + + def test_str_default_length(self): + expected = ("name", np.dtype(f"U{array_module._DEFAULT_STR_LENGTH}")) + assert _parse_annotation_post_25("name", None, (np.str_,), {}) == expected + + def test_unsupported_shape_raises(self): + with pytest.raises(ValueError, match="not understood or supported"): + _parse_annotation_post_25("value", "bad", (), {})