ENH: Add dtype.__class_getitem__

Author: Bas van Beek

numpy/__init__.pyi

@@ -1091,6 +1091,9 @@ class dtype(Generic[_DTypeScalar_co]):
         copy: bool = ...,
     ) -> dtype[object_]: ...
 
+    if sys.version_info >= (3, 9):
+        def __class_getitem__(self, item: Any) -> GenericAlias: ...
+
     @overload
     def __getitem__(self: dtype[void], key: List[str]) -> dtype[void]: ...
     @overload

numpy/core/_add_newdocs.py

@@ -6079,6 +6079,36 @@
 
     """))
 
+if sys.version_info >= (3, 9):
+    add_newdoc('numpy.core.multiarray', 'dtype', ('__class_getitem__',
+        """
+        __class_getitem__(item, /)
+
+        Return a parametrized wrapper around the `~numpy.dtype` type.
+
+        .. versionadded:: 1.22
+
+        Returns
+        -------
+        alias : types.GenericAlias
+            A parametrized `~numpy.dtype` type.
+
+        Examples
+        --------
+        >>> import numpy as np
+
+        >>> np.dtype[np.int64]
+        numpy.dtype[numpy.int64]
+
+        Note
+        ----
+        This method is only available for python 3.9 and later.
+
+        See Also
+        --------
+        :pep:`585` : Type hinting generics in standard collections.
+
+        """))
 
 ##############################################################################
 #

numpy/core/src/multiarray/descriptor.c

@@ -257,7 +257,7 @@ static PyArray_Descr *
 _convert_from_tuple(PyObject *obj, int align)
 {
     if (PyTuple_GET_SIZE(obj) != 2) {
-        PyErr_Format(PyExc_TypeError, 
+        PyErr_Format(PyExc_TypeError,
 	        "Tuple must have size 2, but has size %zd",
 	        PyTuple_GET_SIZE(obj));
         return NULL;
@@ -449,8 +449,8 @@ _convert_from_array_descr(PyObject *obj, int align)
     for (int i = 0; i < n; i++) {
         PyObject *item = PyList_GET_ITEM(obj, i);
         if (!PyTuple_Check(item) || (PyTuple_GET_SIZE(item) < 2)) {
-            PyErr_Format(PyExc_TypeError, 
-			 "Field elements must be 2- or 3-tuples, got '%R'", 
+            PyErr_Format(PyExc_TypeError,
+			 "Field elements must be 2- or 3-tuples, got '%R'",
 			 item);
             goto fail;
         }
@@ -461,7 +461,7 @@ _convert_from_array_descr(PyObject *obj, int align)
         }
         else if (PyTuple_Check(name)) {
             if (PyTuple_GET_SIZE(name) != 2) {
-                PyErr_Format(PyExc_TypeError, 
+                PyErr_Format(PyExc_TypeError,
 				"If a tuple, the first element of a field tuple must have "
 				"two elements, not %zd",
 			       	PyTuple_GET_SIZE(name));
@@ -475,7 +475,7 @@ _convert_from_array_descr(PyObject *obj, int align)
             }
         }
         else {
-            PyErr_SetString(PyExc_TypeError, 
+            PyErr_SetString(PyExc_TypeError,
 			            "First element of field tuple is "
 			            "neither a tuple nor str");
             goto fail;
@@ -3112,6 +3112,13 @@ static PyMethodDef arraydescr_methods[] = {
     {"newbyteorder",
         (PyCFunction)arraydescr_newbyteorder,
         METH_VARARGS, NULL},
+
+    /* for typing; requires python >= 3.9 */
+    #ifdef Py_GENERICALIASOBJECT_H
+    {"__class_getitem__",
+        (PyCFunction)Py_GenericAlias,
+        METH_CLASS | METH_O, NULL},
+    #endif
     {NULL, NULL, 0, NULL}           /* sentinel */
 };
 

numpy/core/tests/test_dtype.py

@@ -4,6 +4,8 @@
 import ctypes
 import gc
 import warnings
+import types
+from typing import Any
 
 import numpy as np
 from numpy.core._rational_tests import rational
@@ -111,9 +113,9 @@ def test_richcompare_invalid_dtype_comparison(self, operation):
     @pytest.mark.parametrize("dtype",
              ['Bool', 'Bytes0', 'Complex32', 'Complex64',
               'Datetime64', 'Float16', 'Float32', 'Float64',
-              'Int8', 'Int16', 'Int32', 'Int64', 
+              'Int8', 'Int16', 'Int32', 'Int64',
               'Object0', 'Str0', 'Timedelta64',
-              'UInt8', 'UInt16', 'Uint32', 'UInt32', 
+              'UInt8', 'UInt16', 'Uint32', 'UInt32',
               'Uint64', 'UInt64', 'Void0',
               "Float128", "Complex128"])
     def test_numeric_style_types_are_invalid(self, dtype):
@@ -1549,3 +1551,18 @@ class mytype:
             # Tests that a dtype must have its type field set up to np.dtype
             # or in this case a builtin instance.
             create_custom_field_dtype(blueprint, mytype, 2)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 9), reason="Requires python 3.9")
+class TestClassGetItem:
+    def test_dtype(self) -> None:
+        alias = np.dtype[Any]
+        assert isinstance(alias, types.GenericAlias)
+        assert alias.__origin__ is np.dtype
+
+    @pytest.mark.parametrize("code", np.typecodes["All"])
+    def test_dtype_subclass(self, code: str) -> None:
+        cls = type(np.dtype(code))
+        alias = cls[Any]
+        assert isinstance(alias, types.GenericAlias)
+        assert alias.__origin__ is cls