ENH: Add ndarray.__class_getitem__

Author: Bas van Beek

numpy/__init__.pyi

@@ -9,6 +9,9 @@ from abc import abstractmethod
 from types import TracebackType, MappingProxyType
 from contextlib import ContextDecorator
 
+if sys.version_info >= (3, 9):
+    from types import GenericAlias
+
 from numpy._pytesttester import PytestTester
 from numpy.core.multiarray import flagsobj
 from numpy.core._internal import _ctypes
@@ -1697,6 +1700,10 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
         strides: _ShapeLike = ...,
         order: _OrderKACF = ...,
     ) -> _ArraySelf: ...
+
+    if sys.version_info >= (3, 9):
+        def __class_getitem__(self, item: Any) -> GenericAlias: ...
+
     @overload
     def __array__(self, dtype: None = ..., /) -> ndarray[Any, _DType_co]: ...
     @overload

numpy/core/_add_newdocs.py

@@ -9,6 +9,7 @@
 
 """
 
+import sys
 from numpy.core.function_base import add_newdoc
 from numpy.core.overrides import array_function_like_doc
 
@@ -796,7 +797,7 @@
     object : array_like
         An array, any object exposing the array interface, an object whose
         __array__ method returns an array, or any (nested) sequence.
-        If object is a scalar, a 0-dimensional array containing object is 
+        If object is a scalar, a 0-dimensional array containing object is
         returned.
     dtype : data-type, optional
         The desired data-type for the array.  If not given, then the type will
@@ -2201,8 +2202,8 @@
     empty : Create an array, but leave its allocated memory unchanged (i.e.,
             it contains "garbage").
     dtype : Create a data-type.
-    numpy.typing.NDArray : A :term:`generic <generic type>` version
-                           of ndarray.
+    numpy.typing.NDArray : An ndarray alias :term:`generic <generic type>`
+                           w.r.t. its `dtype.type <numpy.dtype.type>`.
 
     Notes
     -----
@@ -2798,6 +2799,40 @@
     """))
 
 
+if sys.version_info > (3, 9):
+    add_newdoc('numpy.core.multiarray', 'ndarray', ('__class_getitem__',
+        """a.__class_getitem__(item, /)
+
+        Return a parametrized wrapper around the `~numpy.ndarray` type.
+
+        .. versionadded:: 1.22
+
+        Returns
+        -------
+        alias : types.GenericAlias
+            A parametrized `~numpy.ndarray` type.
+
+        Examples
+        --------
+        >>> from typing import Any
+        >>> import numpy as np
+
+        >>> np.ndarray[Any, np.dtype]
+        numpy.ndarray[typing.Any, numpy.dtype]
+
+        Note
+        ----
+        This method is only available for python 3.9 and later.
+
+        See Also
+        --------
+        :pep:`585` : Type hinting generics in standard collections.
+        numpy.typing.NDArray : An ndarray alias :term:`generic <generic type>`
+                            w.r.t. its `dtype.type <numpy.dtype.type>`.
+
+        """))
+
+
 add_newdoc('numpy.core.multiarray', 'ndarray', ('__deepcopy__',
     """a.__deepcopy__(memo, /) -> Deep copy of array.
 

numpy/core/src/multiarray/methods.c

@@ -2756,6 +2756,13 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = {
         (PyCFunction) array_format,
         METH_VARARGS, NULL},
 
+    /* for typing; requires python >= 3.9 */
+    #ifdef Py_GENERICALIASOBJECT_H
+    {"__class_getitem__",
+        (PyCFunction)Py_GenericAlias,
+        METH_CLASS | METH_O, NULL},
+    #endif
+
     /* Original and Extended methods added 2005 */
     {"all",
         (PyCFunction)array_all,

numpy/core/tests/test_arraymethod.py

@@ -3,6 +3,10 @@
 this is private API, but when added, public API may be added here.
 """
 
+import sys
+import types
+from typing import Any, Type
+
 import pytest
 
 import numpy as np
@@ -56,3 +60,14 @@ def test_invalid_arguments(self, args, error):
         # This is private API, which may be modified freely
         with pytest.raises(error):
             self.method._simple_strided_call(*args)
+
+
+@pytest.mark.parametrize(
+    "cls", [np.ndarray, np.recarray, np.chararray, np.matrix, np.memmap]
+)
+@pytest.mark.skipif(sys.version_info < (3, 9), reason="Requires python 3.9")
+def test_class_getitem(cls: Type[np.ndarray]) -> None:
+    """Test `ndarray.__class_getitem__`."""
+    alias = cls[Any, Any]
+    assert isinstance(alias, types.GenericAlias)
+    assert alias.__origin__ is cls