ENH: Add special-casing for complexfloating so that it can take 2 parameters

Author: BvB93

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

@@ -1812,12 +1812,22 @@ numbertype_class_getitem_abc(PyObject *cls, PyObject *args)
 
 #ifdef Py_GENERICALIASOBJECT_H
     Py_ssize_t args_len;
+    int args_len_expected;
+
+    /* complexfloating should take 2 parameters, all others take 1 */
+    if (PyType_IsSubtype((PyTypeObject *)cls,
+                         &PyComplexFloatingArrType_Type)) {
+        args_len_expected = 2;
+    }
+    else {
+        args_len_expected = 1;
+    }
 
     args_len = PyTuple_Check(args) ? PyTuple_Size(args) : 1;
-    if (args_len != 1) {
+    if (args_len != args_len_expected) {
         return PyErr_Format(PyExc_TypeError,
                             "Too %s arguments for %s",
-                            args_len > 1 ? "many" : "few",
+                            args_len > args_len_expected ? "many" : "few",
                             ((PyTypeObject *)cls)->tp_name);
     }
     generic_alias = Py_GenericAlias(cls, args);

numpy/core/tests/test_scalar_methods.py

@@ -142,13 +142,17 @@ class TestClassGetItem:
         np.unsignedinteger,
         np.signedinteger,
         np.floating,
-        np.complexfloating,
     ])
     def test_abc(self, cls: Type[np.number]) -> None:
         alias = cls[Any]
         assert isinstance(alias, types.GenericAlias)
         assert alias.__origin__ is cls
 
+    def test_abc_complexfloating(self) -> None:
+        alias = np.complexfloating[Any, Any]
+        assert isinstance(alias, types.GenericAlias)
+        assert alias.__origin__ is np.complexfloating
+
     @pytest.mark.parametrize("cls", [np.generic, np.flexible, np.character])
     def test_abc_non_numeric(self, cls: Type[np.generic]) -> None:
         with pytest.raises(TypeError):
@@ -174,7 +178,7 @@ def test_subscript_scalar(self) -> None:
 
 
 @pytest.mark.skipif(sys.version_info >= (3, 9), reason="Requires python 3.8")
-@pytest.mark.parametrize("cls", [np.number, np.int64])
+@pytest.mark.parametrize("cls", [np.number, np.complexfloating, np.int64])
 def test_class_getitem_38(cls: Type[np.number]) -> None:
     match = "Type subscription requires python >= 3.9"
     with pytest.raises(TypeError, match=match):