chore: cleanup of get_dtype (with numpy 2.5 support)

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 @@
     @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:

tests/unit/model/arrays/test_get_dtype.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
+#
+# 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", (), {})