numpy: add Array4D and its 4-d siblings

docs/reference/numpy/array-likes.md

@@ -309,6 +309,10 @@ matrix-likes, and cuboid-likes, and the `To{}` aliases for "bare" scalar types.
     wasn't generic, making it impossible to distinguish between `np.False_` and `np.True_`
     using static typing.
 
+!!! note
+
+    For 4-d array-likes there are only the generic `ToArray4D` and `ToArrayStrict4D`.
+
 !!! note
 
     The `ToArrayStrict{1,2,3}D` types are generic since `optype 0.9.1`, analogous to

docs/reference/numpy/shape.md

@@ -44,9 +44,9 @@ their differences become clear:
     </tr>
 </table>
 
-Additionally, there are the four `Array{0,1,2,3}D` aliases, which are
-equivalent to `Array` with `tuple[()]`, `tuple[int]`, `tuple[int, int]` and
-`tuple[int, int, int]` as shape-type, respectively.
+Additionally, there are the five `Array{0,1,2,3,4}D` aliases, which are
+equivalent to `Array` with `tuple[()]`, `tuple[int]`, `tuple[int, int]`,
+`tuple[int, int, int]` and `tuple[int, int, int, int]` as shape-type, respectively.
 
 !!! info
 
@@ -102,12 +102,12 @@ In the same way as `ArrayND` for `ndarray` (shown for reference), its subtypes
     </tr>
 </table>
 
-For masked arrays with specific `ndim`, you could also use one of the four
-`MArray{0,1,2,3}D` aliases.
+For masked arrays with specific `ndim`, you could also use one of the five
+`MArray{0,1,2,3,4}D` aliases.
 
 ## Array typeguards
 
-To check whether a given object is an instance of `Array{0,1,2,3,N}D`, in a way that
+To check whether a given object is an instance of `Array{0,1,2,3,4,N}D`, in a way that
 static type-checkers also understand it, the following [PEP 742][PEP742] typeguards can
 be used:
 
@@ -146,6 +146,11 @@ be used:
         <td><code>Array3D[ST]</code></td>
         <td><code>tuple[int, int, int]</code></td>
     </tr>
+    <tr>
+        <td><code>is_array_4d</code></td>
+        <td><code>Array4D[ST]</code></td>
+        <td><code>tuple[int, int, int, int]</code></td>
+    </tr>
 </table>
 
 These functions additionally accept an optional `dtype` argument, that can either be

optype/numpy/_array.py

@@ -13,11 +13,11 @@
 from optype._utils import set_module
 
 __all__ = [  # ruff: ignore[unsorted-dunder-all]
-    "Array", "Array0D", "Array1D", "Array2D", "Array3D", "ArrayND",
-    "MArray", "MArray0D", "MArray1D", "MArray2D", "MArray3D",
+    "Array", "Array0D", "Array1D", "Array2D", "Array3D", "Array4D", "ArrayND",
+    "MArray", "MArray0D", "MArray1D", "MArray2D", "MArray3D", "MArray4D",
     "Matrix",
     "CanArray", "CanArray0",
-    "CanArrayND", "CanArray0D", "CanArray1D", "CanArray2D", "CanArray3D",
+    "CanArrayND", "CanArray0D", "CanArray1D", "CanArray2D", "CanArray3D", "CanArray4D",
     "CanArrayFinalize", "CanArrayWrap",
     "HasArrayInterface", "HasArrayPriority",
 ]  # fmt: skip
@@ -148,6 +148,11 @@ def __array__(self, /) -> np.ndarray[_NDT_co, np.dtype[_SCT0_co]]: ...
     np.ndarray[tuple[int, int, int], np.dtype[_SCT]],
     type_params=(_SCT,),
 )
+Array4D = TypeAliasType(
+    "Array4D",
+    np.ndarray[tuple[int, int, int, int], np.dtype[_SCT]],
+    type_params=(_SCT,),
+)
 
 MArray0D = TypeAliasType(
     "MArray0D",
@@ -169,6 +174,11 @@ def __array__(self, /) -> np.ndarray[_NDT_co, np.dtype[_SCT0_co]]: ...
     np.ma.MaskedArray[tuple[int, int, int], np.dtype[_SCT]],
     type_params=(_SCT,),
 )
+MArray4D = TypeAliasType(
+    "MArray4D",
+    np.ma.MaskedArray[tuple[int, int, int, int], np.dtype[_SCT]],
+    type_params=(_SCT,),
+)
 
 
 ###########################
@@ -233,6 +243,18 @@ def __len__(self, /) -> int: ...
     def __array__(self, /) -> np.ndarray[tuple[int, int, int], np.dtype[_SCT_co]]: ...
 
 
+@runtime_checkable
+@set_module("optype.numpy")
+class CanArray4D(Protocol[_SCT_co]):
+    """The 4-d variant of `optype.numpy.CanArrayND`."""
+
+    def __len__(self, /) -> int: ...
+    def __array__(
+        self,
+        /,
+    ) -> np.ndarray[tuple[int, int, int, int], np.dtype[_SCT_co]]: ...
+
+
 # this is almost always a `ndarray`, but setting a `bound` might break in some
 # edge cases
 _T_contra = TypeVar("_T_contra", contravariant=True, default=Any)

optype/numpy/_is.py

@@ -8,14 +8,15 @@
 
 import numpy as np
 
-from ._array import Array0D, Array1D, Array2D, Array3D, ArrayND
+from ._array import Array0D, Array1D, Array2D, Array3D, Array4D, ArrayND
 from ._dtype import ToDType
 
 __all__ = [
     "is_array_0d",
     "is_array_1d",
     "is_array_2d",
     "is_array_3d",
+    "is_array_4d",
     "is_array_nd",
     "is_dtype",
     "is_sctype",
@@ -100,3 +101,12 @@ def is_array_3d(
 ) -> TypeIs[Array3D[ScalarT]]:
     """Checks if `a` is a 3-d `ndarray` of the given dtype (defaults to `generic`)."""
     return is_array_nd(a, dtype) and a.ndim == 3
+
+
+def is_array_4d(
+    a: Any,
+    /,
+    dtype: ToDType[ScalarT] | None = None,
+) -> TypeIs[Array4D[ScalarT]]:
+    """Checks if `a` is a 4-d `ndarray` of the given dtype (defaults to `generic`)."""
+    return is_array_nd(a, dtype) and a.ndim == 4

optype/numpy/_to.py

@@ -27,6 +27,7 @@
     "ToArray1D", "ToArrayStrict1D",
     "ToArray2D", "ToArrayStrict2D",
     "ToArray3D", "ToArrayStrict3D",
+    "ToArray4D", "ToArrayStrict4D",
     "ToArrayND",
 
     "ToFalse", "ToTrue",
@@ -159,6 +160,11 @@ def __array__(self, /) -> np.ndarray[Any, np.dtype[SCT_co]]: ...
         nptc.CanArray[tuple[int, int, int], np.dtype[SCT]],
         type_params=(SCT,),
     )
+    _CanArrayStrict4D = TypeAliasType(
+        "_CanArrayStrict4D",
+        nptc.CanArray[tuple[int, int, int, int], np.dtype[SCT]],
+        type_params=(SCT,),
+    )
 
 else:
 
@@ -174,6 +180,12 @@ def __array__(self) -> np.ndarray[tuple[int, int], np.dtype[SCT_co]]: ...
     class _CanArrayStrict3D(Protocol[SCT_co]):
         def __array__(self) -> np.ndarray[tuple[int, int, int], np.dtype[SCT_co]]: ...
 
+    @runtime_checkable
+    class _CanArrayStrict4D(Protocol[SCT_co]):
+        def __array__(
+            self,
+        ) -> np.ndarray[tuple[int, int, int, int], np.dtype[SCT_co]]: ...
+
 
 type _To1D1[SCT: np.generic] = _CanArrayND[SCT] | Seq[SCT]
 type _To1D2[T, SCT: np.generic] = _CanArrayND[SCT] | Seq[T | SCT]
@@ -184,6 +196,8 @@ def __array__(self) -> np.ndarray[tuple[int, int, int], np.dtype[SCT_co]]: ...
 type _To3D1[SCT: np.generic] = _CanArrayND[SCT] | Seq[_To2D1[SCT]]
 type _To3D2[T, SCT: np.generic] = _CanArrayND[SCT] | Seq[_To2D2[T, SCT]]
 
+type _To4D2[T, SCT: np.generic] = _CanArrayND[SCT] | Seq[_To3D2[T, SCT]]
+
 type _ToND1[SCT: np.generic] = _CanArrayND[SCT] | SeqND[_CanArray[SCT]]
 type _ToND2[T, SCT: np.generic] = _CanArrayND[SCT] | SeqND[T | _CanArray[SCT]]
 
@@ -200,6 +214,10 @@ def __array__(self) -> np.ndarray[tuple[int, int, int], np.dtype[SCT_co]]: ...
     _CanArrayStrict3D[SCT] | Seq[_ToStrict2D2[T, SCT]]
 )
 
+type _ToStrict4D2[T, SCT: np.generic] = (
+    _CanArrayStrict4D[SCT] | Seq[_ToStrict3D2[T, SCT]]
+)
+
 
 ###
 
@@ -237,6 +255,7 @@ def __array__(self) -> np.ndarray[tuple[int, int, int], np.dtype[SCT_co]]: ...
 ToArray1D = TypeAliasType("ToArray1D", _To1D2[T, SCT], type_params=(T, SCT))
 ToArray2D = TypeAliasType("ToArray2D", _To2D2[T, SCT], type_params=(T, SCT))
 ToArray3D = TypeAliasType("ToArray3D", _To3D2[T, SCT], type_params=(T, SCT))
+ToArray4D = TypeAliasType("ToArray4D", _To4D2[T, SCT], type_params=(T, SCT))
 ToArrayND = TypeAliasType("ToArrayND", _ToND2[T, SCT], type_params=(T, SCT))
 
 type ToFalse = nptc.LiteralFalse | Literal[0]
@@ -390,6 +409,11 @@ def __array__(self) -> np.ndarray[tuple[int, int, int], np.dtype[SCT_co]]: ...
     _ToStrict3D2[T, SCT],
     type_params=(T, SCT),
 )
+ToArrayStrict4D = TypeAliasType(
+    "ToArrayStrict4D",
+    _ToStrict4D2[T, SCT],
+    type_params=(T, SCT),
+)
 
 type ToBoolStrict1D = _ToStrict1D2[_PyBool, np.bool]
 type ToBoolStrict2D = _ToStrict2D2[_PyBool, np.bool]

tests/numpy/test_is.py

@@ -51,7 +51,7 @@ def test_is_sctype(sctype: type[np.generic]) -> None:
 )
 @pytest.mark.parametrize("dtype", DTYPES, ids=str)
 @pytest.mark.parametrize("ndtype", NDARRAY_TYPES)
-def test_is_array(
+def test_is_array(  # ruff: ignore[too-many-statements]
     ndtype: type[onp.Array],
     dtype: onp.DType,
     dtype_map: Callable[[onp.DType], onp.DType | type[np.generic]],
@@ -63,47 +63,55 @@ def test_is_array(
     assert not onp.is_array_1d(arr[0])
     assert not onp.is_array_2d(arr[0])
     assert not onp.is_array_3d(arr[0])
+    assert not onp.is_array_4d(arr[0])
 
     assert onp.is_array_nd(arr[1])
     assert not onp.is_array_0d(arr[1])
     assert onp.is_array_1d(arr[1])
     assert not onp.is_array_2d(arr[1])
     assert not onp.is_array_3d(arr[1])
+    assert not onp.is_array_4d(arr[1])
 
     assert onp.is_array_nd(arr[2])
     assert not onp.is_array_0d(arr[2])
     assert not onp.is_array_1d(arr[2])
     assert onp.is_array_2d(arr[2])
     assert not onp.is_array_3d(arr[2])
+    assert not onp.is_array_4d(arr[2])
 
     assert onp.is_array_nd(arr[2].view(np.matrix))
     assert not onp.is_array_0d(arr[2].view(np.matrix))
     assert not onp.is_array_1d(arr[2].view(np.matrix))
     assert onp.is_array_2d(arr[2].view(np.matrix))
     assert not onp.is_array_3d(arr[2].view(np.matrix))
+    assert not onp.is_array_4d(arr[2].view(np.matrix))
 
     assert onp.is_array_nd(arr[3])
     assert not onp.is_array_0d(arr[3])
     assert not onp.is_array_1d(arr[3])
     assert not onp.is_array_2d(arr[3])
     assert onp.is_array_3d(arr[3])
+    assert not onp.is_array_4d(arr[3])
 
     assert onp.is_array_nd(arr[4])
     assert not onp.is_array_0d(arr[4])
     assert not onp.is_array_1d(arr[4])
     assert not onp.is_array_2d(arr[4])
     assert not onp.is_array_3d(arr[4])
+    assert onp.is_array_4d(arr[4])
 
     dtype_is = dtype_map(dtype)
     assert onp.is_array_0d(arr[0], dtype=dtype_is)
     assert onp.is_array_1d(arr[1], dtype=dtype_is)
     assert onp.is_array_2d(arr[2], dtype=dtype_is)
     assert onp.is_array_3d(arr[3], dtype=dtype_is)
+    assert onp.is_array_4d(arr[4], dtype=dtype_is)
     assert onp.is_array_nd(arr[4], dtype=dtype_is)
 
     dtype_not = np.dtype("?") if dtype.char != "?" else np.dtype("B")
     assert not onp.is_array_0d(arr[0], dtype=dtype_not)
     assert not onp.is_array_1d(arr[1], dtype=dtype_not)
     assert not onp.is_array_2d(arr[2], dtype=dtype_not)
     assert not onp.is_array_3d(arr[3], dtype=dtype_not)
+    assert not onp.is_array_4d(arr[4], dtype=dtype_not)
     assert not onp.is_array_nd(arr[4], dtype=dtype_not)

tests/numpy/test_to.pyi

@@ -101,6 +101,8 @@ type _Arr3_c_co = onp.Array3D[_Sca_c_co] | Seq[_Arr2_c_co]
 type _Arr3_f8_co = onp.Array3D[_Sca_f8_co] | Seq[_Arr2_f8_co]
 type _Arr3_c16_co = onp.Array3D[_Sca_c16_co] | Seq[_Arr2_c16_co]
 
+type _Arr4_x = onp.Array4D[_Sca_x] | Seq[_Arr3_x]
+
 x_: _Val_x
 b_: _Val_b
 i_: _Val_i
@@ -166,6 +168,8 @@ c_co_3d: _Arr3_c_co
 f8_co_3d: _Arr3_f8_co
 c16_co_3d: _Arr3_c16_co
 
+x_4d: _Arr4_x
+
 # scalar
 
 def sca_sca() -> None:
@@ -358,6 +362,19 @@ def s3d_a3d() -> None:
     f8_co__f: onp.ToFloat64Strict3D = f_co_3d  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
     c16_co__c: onp.ToComplex128Strict3D = c_co_3d  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
 
+# 4-d
+
+def a4d() -> None:
+    x__sca: onp.ToArray4D = x_  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
+    x__a3d: onp.ToArray4D = x_3d  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
+    x__a4d: onp.ToArray4D = x_4d
+    x__s4d: onp.ToArrayStrict4D = x_4d
+
+def nd_4d() -> None:
+    to_x_4d: onp.ToArray4D
+
+    x__x: onp.ToArrayND = to_x_4d
+
 # n-d
 
 def nd_sca() -> None: