MAINT: Make __class_getitem__ available to all python version and perform basic validation of its input arguments

It will still raise on python 3.8, but now with a more explicit exception message

Author: Bas van Beek

numpy/core/_add_newdocs.py

@@ -9,7 +9,6 @@
 
 """
 
-import sys
 from numpy.core.function_base import add_newdoc
 from numpy.core.overrides import array_function_like_doc
 
@@ -2799,38 +2798,37 @@
     """))
 
 
-if sys.version_info > (3, 9):
-    add_newdoc('numpy.core.multiarray', 'ndarray', ('__class_getitem__',
-        """a.__class_getitem__(item, /)
+add_newdoc('numpy.core.multiarray', 'ndarray', ('__class_getitem__',
+    """a.__class_getitem__(item, /)
 
-        Return a parametrized wrapper around the `~numpy.ndarray` type.
+    Return a parametrized wrapper around the `~numpy.ndarray` type.
 
-        .. versionadded:: 1.22
+    .. versionadded:: 1.22
 
-        Returns
-        -------
-        alias : types.GenericAlias
-            A parametrized `~numpy.ndarray` type.
+    Returns
+    -------
+    alias : types.GenericAlias
+        A parametrized `~numpy.ndarray` type.
 
-        Examples
-        --------
-        >>> from typing import Any
-        >>> import numpy as np
+    Examples
+    --------
+    >>> from typing import Any
+    >>> import numpy as np
 
-        >>> np.ndarray[Any, np.dtype]
-        numpy.ndarray[typing.Any, numpy.dtype]
+    >>> np.ndarray[Any, np.dtype[Any]]
+    numpy.ndarray[typing.Any, numpy.dtype[Any]]
 
-        Note
-        ----
-        This method is only available for python 3.9 and later.
+    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>`.
+    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__',
@@ -6079,36 +6077,35 @@
 
     """))
 
-if sys.version_info >= (3, 9):
-    add_newdoc('numpy.core.multiarray', 'dtype', ('__class_getitem__',
-        """
-        __class_getitem__(item, /)
+add_newdoc('numpy.core.multiarray', 'dtype', ('__class_getitem__',
+    """
+    __class_getitem__(item, /)
 
-        Return a parametrized wrapper around the `~numpy.dtype` type.
+    Return a parametrized wrapper around the `~numpy.dtype` type.
 
-        .. versionadded:: 1.22
+    .. versionadded:: 1.22
 
-        Returns
-        -------
-        alias : types.GenericAlias
-            A parametrized `~numpy.dtype` type.
+    Returns
+    -------
+    alias : types.GenericAlias
+        A parametrized `~numpy.dtype` type.
 
-        Examples
-        --------
-        >>> import numpy as np
+    Examples
+    --------
+    >>> import numpy as np
 
-        >>> np.dtype[np.int64]
-        numpy.dtype[numpy.int64]
+    >>> np.dtype[np.int64]
+    numpy.dtype[numpy.int64]
 
-        Note
-        ----
-        This method is only available for python 3.9 and later.
+    Note
+    ----
+    This method is only available for python 3.9 and later.
 
-        See Also
-        --------
-        :pep:`585` : Type hinting generics in standard collections.
+    See Also
+    --------
+    :pep:`585` : Type hinting generics in standard collections.
 
-        """))
+    """))
 
 ##############################################################################
 #
@@ -6530,37 +6527,36 @@ def refer_to_array_attribute(attr, method=True):
 add_newdoc('numpy.core.numerictypes', 'generic',
            refer_to_array_attribute('view'))
 
-if sys.version_info >= (3, 9):
-    add_newdoc('numpy.core.numerictypes', 'number', ('__class_getitem__',
-        """
-        __class_getitem__(item, /)
+add_newdoc('numpy.core.numerictypes', 'number', ('__class_getitem__',
+    """
+    __class_getitem__(item, /)
 
-        Return a parametrized wrapper around the `~numpy.number` type.
+    Return a parametrized wrapper around the `~numpy.number` type.
 
-        .. versionadded:: 1.22
+    .. versionadded:: 1.22
 
-        Returns
-        -------
-        alias : types.GenericAlias
-            A parametrized `~numpy.number` type.
+    Returns
+    -------
+    alias : types.GenericAlias
+        A parametrized `~numpy.number` type.
 
-        Examples
-        --------
-        >>> from typing import Any
-        >>> import numpy as np
+    Examples
+    --------
+    >>> from typing import Any
+    >>> import numpy as np
 
-        >>> np.signedinteger[Any]
-        numpy.signedinteger[typing.Any]
+    >>> np.signedinteger[Any]
+    numpy.signedinteger[typing.Any]
 
-        Note
-        ----
-        This method is only available for python 3.9 and later.
+    Note
+    ----
+    This method is only available for python 3.9 and later.
 
-        See Also
-        --------
-        :pep:`585` : Type hinting generics in standard collections.
+    See Also
+    --------
+    :pep:`585` : Type hinting generics in standard collections.
 
-        """))
+    """))
 
 ##############################################################################
 #

numpy/core/src/multiarray/descriptor.c

@@ -3101,6 +3101,30 @@ arraydescr_newbyteorder(PyArray_Descr *self, PyObject *args)
     return (PyObject *)PyArray_DescrNewByteorder(self, endian);
 }
 
+static PyObject *
+arraydescr_class_getitem(PyObject *cls, PyObject *args)
+{
+    PyObject *generic_alias;
+
+#ifdef Py_GENERICALIASOBJECT_H
+    Py_ssize_t args_len;
+
+    args_len = PyTuple_Check(args) ? PyTuple_Size(args) : 1;
+    if (args_len != 1) {
+        return PyErr_Format(PyExc_TypeError,
+                            "Too %s arguments for %s",
+                            args_len > 1 ? "many" : "few",
+                            ((PyTypeObject *)cls)->tp_name);
+    }
+    generic_alias = Py_GenericAlias(cls, args);
+#else
+    PyErr_SetString(PyExc_TypeError,
+                    "Type subscription requires python >= 3.9");
+    generic_alias = NULL;
+#endif
+    return generic_alias;
+}
+
 static PyMethodDef arraydescr_methods[] = {
     /* for pickling */
     {"__reduce__",
@@ -3112,13 +3136,10 @@ 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,
+        (PyCFunction)arraydescr_class_getitem,
         METH_CLASS | METH_O, NULL},
-    #endif
     {NULL, NULL, 0, NULL}           /* sentinel */
 };
 

numpy/core/src/multiarray/methods.c

@@ -2699,6 +2699,30 @@ array_complex(PyArrayObject *self, PyObject *NPY_UNUSED(args))
     return c;
 }
 
+static PyObject *
+array_class_getitem(PyObject *cls, PyObject *args)
+{
+    PyObject *generic_alias;
+
+#ifdef Py_GENERICALIASOBJECT_H
+    Py_ssize_t args_len;
+
+    args_len = PyTuple_Check(args) ? PyTuple_Size(args) : 1;
+    if (args_len != 2) {
+        return PyErr_Format(PyExc_TypeError,
+                            "Too %s arguments for %s",
+                            args_len > 2 ? "many" : "few",
+                            ((PyTypeObject *)cls)->tp_name);
+    }
+    generic_alias = Py_GenericAlias(cls, args);
+#else
+    PyErr_SetString(PyExc_TypeError,
+                    "Type subscription requires python >= 3.9");
+    generic_alias = NULL;
+#endif
+    return generic_alias;
+}
+
 NPY_NO_EXPORT PyMethodDef array_methods[] = {
 
     /* for subtypes */
@@ -2757,11 +2781,9 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = {
         METH_VARARGS, NULL},
 
     /* for typing; requires python >= 3.9 */
-    #ifdef Py_GENERICALIASOBJECT_H
     {"__class_getitem__",
-        (PyCFunction)Py_GenericAlias,
+        (PyCFunction)array_class_getitem,
         METH_CLASS | METH_O, NULL},
-    #endif
 
     /* Original and Extended methods added 2005 */
     {"all",

numpy/core/src/multiarray/scalartypes.c.src

@@ -1805,20 +1805,48 @@ gentype_setflags(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args),
     Py_RETURN_NONE;
 }
 
+static PyObject *
+numbertype_class_getitem_abc(PyObject *cls, PyObject *args)
+{
+    PyObject *generic_alias;
+
+#ifdef Py_GENERICALIASOBJECT_H
+    Py_ssize_t args_len;
+
+    args_len = PyTuple_Check(args) ? PyTuple_Size(args) : 1;
+    if (args_len != 1) {
+        return PyErr_Format(PyExc_TypeError,
+                            "Too %s arguments for %s",
+                            args_len > 1 ? "many" : "few",
+                            ((PyTypeObject *)cls)->tp_name);
+    }
+    generic_alias = Py_GenericAlias(cls, args);
+#else
+    PyErr_SetString(PyExc_TypeError,
+                    "Type subscription requires python >= 3.9");
+    generic_alias = NULL;
+#endif
+    return generic_alias;
+}
+
 /*
  * Use for concrete np.number subclasses, making them act as if they
  * were subtyped from e.g. np.signedinteger[object], thus lacking any
  * free subscription parameters. Requires python >= 3.9.
  */
-#ifdef Py_GENERICALIASOBJECT_H
 static PyObject *
 numbertype_class_getitem(PyObject *cls, PyObject *args)
 {
-    return PyErr_Format(PyExc_TypeError,
-                        "There are no type variables left in %s",
-                        ((PyTypeObject *)cls)->tp_name);
-}
+#ifdef Py_GENERICALIASOBJECT_H
+    PyErr_Format(PyExc_TypeError,
+                 "There are no type variables left in %s",
+                 ((PyTypeObject *)cls)->tp_name);
+#else
+    PyErr_SetString(PyExc_TypeError,
+                    "Type subscription requires python >= 3.9");
 #endif
+    return NULL;
+}
 
 /*
  * casting complex numbers (that don't inherit from Python complex)
@@ -2205,11 +2233,9 @@ static PyGetSetDef inttype_getsets[] = {
 
 static PyMethodDef numbertype_methods[] = {
     /* for typing; requires python >= 3.9 */
-    #ifdef Py_GENERICALIASOBJECT_H
     {"__class_getitem__",
-        (PyCFunction)Py_GenericAlias,
+        (PyCFunction)numbertype_class_getitem_abc,
         METH_CLASS | METH_O, NULL},
-    #endif
     {NULL, NULL, 0, NULL}  /* sentinel */
 };
 
@@ -2221,11 +2247,9 @@ static PyMethodDef @name@type_methods[] = {
         (PyCFunction)@name@_complex,
         METH_VARARGS | METH_KEYWORDS, NULL},
     /* for typing; requires python >= 3.9 */
-    #ifdef Py_GENERICALIASOBJECT_H
     {"__class_getitem__",
         (PyCFunction)numbertype_class_getitem,
         METH_CLASS | METH_O, NULL},
-    #endif
     {NULL, NULL, 0, NULL}
 };
 /**end repeat**/
@@ -2264,11 +2288,9 @@ static PyMethodDef @name@type_methods[] = {
         (PyCFunction)@name@_is_integer,
         METH_NOARGS, NULL},
     /* for typing; requires python >= 3.9 */
-    #ifdef Py_GENERICALIASOBJECT_H
     {"__class_getitem__",
         (PyCFunction)numbertype_class_getitem,
         METH_CLASS | METH_O, NULL},
-    #endif
     {NULL, NULL, 0, NULL}
 };
 /**end repeat**/
@@ -2279,11 +2301,9 @@ static PyMethodDef @name@type_methods[] = {
  */
 static PyMethodDef @name@type_methods[] = {
     /* for typing; requires python >= 3.9 */
-    #ifdef Py_GENERICALIASOBJECT_H
     {"__class_getitem__",
         (PyCFunction)numbertype_class_getitem,
         METH_CLASS | METH_O, NULL},
-    #endif
     {NULL, NULL, 0, NULL}
 };
 /**end repeat**/