TYP: Make array shape type variable covariant and bound

Fixes numpy#25729

This change allows future changes to the static typing of numpy that modify
or only work with certain numbers of dimensions.  It also applies the change
to subclasses of ndarray and adds tests.  It allows users to statically
type their array shapes with subtypes of tuple (e.g. NamedTuple) and tuples
of int subtypes (e.g. Literal or NewType).

For a discussion of the merits of TypeVarTuple vs a tuple-bound TypeVar, see
the linked PR

Author: Jacob-Stevens-Haas

.gitignore

@@ -1,5 +1,6 @@
 # Editor temporary/working/backup files #
 #########################################
+env/
 .#*
 [#]*#
 *~

doc/release/upcoming_changes/26081.improvement.rst

@@ -0,0 +1,10 @@
+``ndarray`` shape-type parameter is now covariant and bound to ``tuple[int, ...]``
+----------------------------------------------------------------------------------
+Static typing for ``ndarray`` is a long-term effort that continues
+with this change.  It is a generic type with type parameters for
+the shape and the data type.  Previously, the shape type parameter could be
+any value.  This change restricts it to a tuple of ints, as one would expect
+from using ``ndarray.shape``.  Further, the shape-type parameter has been
+changed from invariant to covariant.  This change also applies to subpytes of
+``ndarray``, e.g. ``np.ma.MaskedArray``.  See the `typing docs <https://typing.readthedocs.io/en/latest/reference/generics.html#variance-of-generic-types>`_
+for more information.

doc/release/upcoming_changes/README.rst

@@ -40,7 +40,7 @@ So for example: ``123.new_feature.rst`` would have the content::
     The ``my_new_feature`` option is now available for `my_favorite_function`.
     To use it, write ``np.my_favorite_function(..., my_new_feature=True)``.
 
-``highlight`` is usually formatted as bulled points making the fragment
+``highlight`` is usually formatted as bullet points making the fragment
 ``* This is a highlight``.
 
 Note the use of single-backticks to get an internal link (assuming

numpy/__init__.pyi

@@ -1500,10 +1500,8 @@ _DType = TypeVar("_DType", bound=dtype[Any])
 _DType_co = TypeVar("_DType_co", covariant=True, bound=dtype[Any])
 _FlexDType = TypeVar("_FlexDType", bound=dtype[flexible])
 
-# TODO: Set the `bound` to something more suitable once we
-# have proper shape support
-_ShapeType = TypeVar("_ShapeType", bound=Any)
-_ShapeType2 = TypeVar("_ShapeType2", bound=Any)
+_ShapeType_co = TypeVar("_ShapeType_co", covariant=True, bound=tuple[int, ...])
+_ShapeType2 = TypeVar("_ShapeType2", bound=tuple[int, ...])
 _NumberType = TypeVar("_NumberType", bound=number[Any])
 
 if sys.version_info >= (3, 12):
@@ -1553,7 +1551,7 @@ class _SupportsImag(Protocol[_T_co]):
     @property
     def imag(self) -> _T_co: ...
 
-class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
+class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     __hash__: ClassVar[None]
     @property
     def base(self) -> None | NDArray[Any]: ...
@@ -1563,14 +1561,14 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
     def size(self) -> int: ...
     @property
     def real(
-        self: ndarray[_ShapeType, dtype[_SupportsReal[_ScalarType]]],  # type: ignore[type-var]
-    ) -> ndarray[_ShapeType, _dtype[_ScalarType]]: ...
+        self: ndarray[_ShapeType_co, dtype[_SupportsReal[_ScalarType]]],  # type: ignore[type-var]
+    ) -> ndarray[_ShapeType_co, _dtype[_ScalarType]]: ...
     @real.setter
     def real(self, value: ArrayLike) -> None: ...
     @property
     def imag(
-        self: ndarray[_ShapeType, dtype[_SupportsImag[_ScalarType]]],  # type: ignore[type-var]
-    ) -> ndarray[_ShapeType, _dtype[_ScalarType]]: ...
+        self: ndarray[_ShapeType_co, dtype[_SupportsImag[_ScalarType]]],  # type: ignore[type-var]
+    ) -> ndarray[_ShapeType_co, _dtype[_ScalarType]]: ...
     @imag.setter
     def imag(self, value: ArrayLike) -> None: ...
     def __new__(
@@ -1591,11 +1589,11 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
     @overload
     def __array__(
         self, dtype: None = ..., /, *, copy: None | bool = ...
-    ) -> ndarray[_ShapeType, _DType_co]: ...
+    ) -> ndarray[_ShapeType_co, _DType_co]: ...
     @overload
     def __array__(
         self, dtype: _DType, /, *, copy: None | bool = ...
-    ) -> ndarray[_ShapeType, _DType]: ...
+    ) -> ndarray[_ShapeType_co, _DType]: ...
 
     def __array_ufunc__(
         self,
@@ -1646,12 +1644,12 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
     @overload
     def __getitem__(self: NDArray[void], key: str) -> NDArray[Any]: ...
     @overload
-    def __getitem__(self: NDArray[void], key: list[str]) -> ndarray[_ShapeType, _dtype[void]]: ...
+    def __getitem__(self: NDArray[void], key: list[str]) -> ndarray[_ShapeType_co, _dtype[void]]: ...
 
     @property
     def ctypes(self) -> _ctypes[int]: ...
     @property
-    def shape(self) -> _Shape: ...
+    def shape(self) -> _ShapeType_co: ...
     @shape.setter
     def shape(self, value: _ShapeLike) -> None: ...
     @property
@@ -3786,7 +3784,7 @@ _MemMapModeKind: TypeAlias = L[
     "write", "w+",
 ]
 
-class memmap(ndarray[_ShapeType, _DType_co]):
+class memmap(ndarray[_ShapeType_co, _DType_co]):
     __array_priority__: ClassVar[float]
     filename: str | None
     offset: int
@@ -3824,7 +3822,7 @@ class memmap(ndarray[_ShapeType, _DType_co]):
     def __array_finalize__(self, obj: object) -> None: ...
     def __array_wrap__(
         self,
-        array: memmap[_ShapeType, _DType_co],
+        array: memmap[_ShapeType_co, _DType_co],
         context: None | tuple[ufunc, tuple[Any, ...], int] = ...,
         return_scalar: builtins.bool = ...,
     ) -> Any: ...
@@ -3927,7 +3925,7 @@ class poly1d:
         k: None | _ArrayLikeComplex_co | _ArrayLikeObject_co = ...,
     ) -> poly1d: ...
 
-class matrix(ndarray[_ShapeType, _DType_co]):
+class matrix(ndarray[_ShapeType_co, _DType_co]):
     __array_priority__: ClassVar[float]
     def __new__(
         subtype,
@@ -3963,13 +3961,13 @@ class matrix(ndarray[_ShapeType, _DType_co]):
     @overload
     def __getitem__(self: NDArray[void], key: str, /) -> matrix[Any, dtype[Any]]: ...
     @overload
-    def __getitem__(self: NDArray[void], key: list[str], /) -> matrix[_ShapeType, dtype[void]]: ...
+    def __getitem__(self: NDArray[void], key: list[str], /) -> matrix[_ShapeType_co, dtype[void]]: ...
 
     def __mul__(self, other: ArrayLike, /) -> matrix[Any, Any]: ...
     def __rmul__(self, other: ArrayLike, /) -> matrix[Any, Any]: ...
-    def __imul__(self, other: ArrayLike, /) -> matrix[_ShapeType, _DType_co]: ...
+    def __imul__(self, other: ArrayLike, /) -> matrix[_ShapeType_co, _DType_co]: ...
     def __pow__(self, other: ArrayLike, /) -> matrix[Any, Any]: ...
-    def __ipow__(self, other: ArrayLike, /) -> matrix[_ShapeType, _DType_co]: ...
+    def __ipow__(self, other: ArrayLike, /) -> matrix[_ShapeType_co, _DType_co]: ...
 
     @overload
     def sum(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ...
@@ -4065,14 +4063,14 @@ class matrix(ndarray[_ShapeType, _DType_co]):
     @property
     def I(self) -> matrix[Any, Any]: ...
     @property
-    def A(self) -> ndarray[_ShapeType, _DType_co]: ...
+    def A(self) -> ndarray[_ShapeType_co, _DType_co]: ...
     @property
     def A1(self) -> ndarray[Any, _DType_co]: ...
     @property
     def H(self) -> matrix[Any, _DType_co]: ...
     def getT(self) -> matrix[Any, _DType_co]: ...
     def getI(self) -> matrix[Any, Any]: ...
-    def getA(self) -> ndarray[_ShapeType, _DType_co]: ...
+    def getA(self) -> ndarray[_ShapeType_co, _DType_co]: ...
     def getA1(self) -> ndarray[Any, _DType_co]: ...
     def getH(self) -> matrix[Any, _DType_co]: ...
 

numpy/_core/defchararray.pyi

@@ -16,7 +16,7 @@ from numpy import (
     int_,
     object_,
     _OrderKACF,
-    _ShapeType, 
+    _ShapeType_co, 
     _CharDType, 
     _SupportsBuffer,
 )
@@ -35,7 +35,7 @@ from numpy._core.multiarray import compare_chararrays as compare_chararrays
 _SCT = TypeVar("_SCT", str_, bytes_)
 _CharArray = chararray[Any, dtype[_SCT]]
 
-class chararray(ndarray[_ShapeType, _CharDType]):
+class chararray(ndarray[_ShapeType_co, _CharDType]):
     @overload
     def __new__(
         subtype,
@@ -436,20 +436,20 @@ class chararray(ndarray[_ShapeType, _CharDType]):
     ) -> _CharArray[bytes_]: ...
 
     def zfill(self, width: _ArrayLikeInt_co) -> chararray[Any, _CharDType]: ...
-    def capitalize(self) -> chararray[_ShapeType, _CharDType]: ...
-    def title(self) -> chararray[_ShapeType, _CharDType]: ...
-    def swapcase(self) -> chararray[_ShapeType, _CharDType]: ...
-    def lower(self) -> chararray[_ShapeType, _CharDType]: ...
-    def upper(self) -> chararray[_ShapeType, _CharDType]: ...
-    def isalnum(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
-    def isalpha(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
-    def isdigit(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
-    def islower(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
-    def isspace(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
-    def istitle(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
-    def isupper(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
-    def isnumeric(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
-    def isdecimal(self) -> ndarray[_ShapeType, dtype[np.bool]]: ...
+    def capitalize(self) -> chararray[_ShapeType_co, _CharDType]: ...
+    def title(self) -> chararray[_ShapeType_co, _CharDType]: ...
+    def swapcase(self) -> chararray[_ShapeType_co, _CharDType]: ...
+    def lower(self) -> chararray[_ShapeType_co, _CharDType]: ...
+    def upper(self) -> chararray[_ShapeType_co, _CharDType]: ...
+    def isalnum(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
+    def isalpha(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
+    def isdigit(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
+    def islower(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
+    def isspace(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
+    def istitle(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
+    def isupper(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
+    def isnumeric(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
+    def isdecimal(self) -> ndarray[_ShapeType_co, dtype[np.bool]]: ...
 
 __all__: list[str]
 

numpy/_core/records.pyi

@@ -16,7 +16,7 @@ from numpy import (
     void,
     _ByteOrder,
     _SupportsBuffer,
-    _ShapeType,
+    _ShapeType_co,
     _DType_co,
     _OrderKACF,
 )
@@ -49,7 +49,7 @@ class record(void):
     @overload
     def __getitem__(self, key: list[str]) -> record: ...
 
-class recarray(ndarray[_ShapeType, _DType_co]):
+class recarray(ndarray[_ShapeType_co, _DType_co]):
     # NOTE: While not strictly mandatory, we're demanding here that arguments
     # for the `format_parser`- and `dtype`-based dtype constructors are
     # mutually exclusive
@@ -114,7 +114,7 @@ class recarray(ndarray[_ShapeType, _DType_co]):
     @overload
     def __getitem__(self, indx: str) -> NDArray[Any]: ...
     @overload
-    def __getitem__(self, indx: list[str]) -> recarray[_ShapeType, dtype[record]]: ...
+    def __getitem__(self, indx: list[str]) -> recarray[_ShapeType_co, dtype[record]]: ...
     @overload
     def field(self, attr: int | str, val: None = ...) -> Any: ...
     @overload

numpy/ma/core.pyi

@@ -15,9 +15,7 @@ from numpy import (
     angle as angle
 )
 
-# TODO: Set the `bound` to something more suitable once we
-# have proper shape support
-_ShapeType = TypeVar("_ShapeType", bound=Any)
+_ShapeType_co = TypeVar("_ShapeType_co", bound=tuple[int, ...], covariant=True)
 _DType_co = TypeVar("_DType_co", bound=dtype[Any], covariant=True)
 
 __all__: list[str]
@@ -165,7 +163,7 @@ class MaskedIterator:
     def __setitem__(self, index, value): ...
     def __next__(self): ...
 
-class MaskedArray(ndarray[_ShapeType, _DType_co]):
+class MaskedArray(ndarray[_ShapeType_co, _DType_co]):
     __array_priority__: Any
     def __new__(cls, data=..., mask=..., dtype=..., copy=..., subok=..., ndmin=..., fill_value=..., keep_mask=..., hard_mask=..., shrink=..., order=...): ...
     def __array_finalize__(self, obj): ...
@@ -300,7 +298,7 @@ class MaskedArray(ndarray[_ShapeType, _DType_co]):
     def __reduce__(self): ...
     def __deepcopy__(self, memo=...): ...
 
-class mvoid(MaskedArray[_ShapeType, _DType_co]):
+class mvoid(MaskedArray[_ShapeType_co, _DType_co]):
     def __new__(
         self,
         data,

numpy/ma/mrecords.pyi

@@ -5,12 +5,10 @@ from numpy.ma import MaskedArray
 
 __all__: list[str]
 
-# TODO: Set the `bound` to something more suitable once we
-# have proper shape support
-_ShapeType = TypeVar("_ShapeType", bound=Any)
+_ShapeType_co = TypeVar("_ShapeType_co", covariant=True, bound=tuple[int, ...])
 _DType_co = TypeVar("_DType_co", bound=dtype[Any], covariant=True)
 
-class MaskedRecords(MaskedArray[_ShapeType, _DType_co]):
+class MaskedRecords(MaskedArray[_ShapeType_co, _DType_co]):
     def __new__(
         cls,
         shape,

numpy/typing/tests/data/fail/shape.pyi

@@ -0,0 +1,6 @@
+from typing import Any
+import numpy as np
+
+# test bounds of _ShapeType_co
+
+np.ndarray[tuple[str, str], Any]  # E: Value of type variable

numpy/typing/tests/data/pass/shape.py

@@ -0,0 +1,18 @@
+from typing import Any, NamedTuple
+
+import numpy as np
+from typing_extensions import assert_type
+
+
+# Subtype of tuple[int, int]
+class XYGrid(NamedTuple):
+    x_axis: int
+    y_axis: int
+
+arr: np.ndarray[XYGrid, Any] = np.empty(XYGrid(2, 2))
+
+# Test variance of _ShapeType_co
+def accepts_2d(a: np.ndarray[tuple[int, int], Any]) -> None:
+    return None
+
+accepts_2d(arr)