bpo-32226: Make __class_getitem__ an automatic class method. (#5098)

2 files changed, 17 insertions, 4 deletions

diff --git a/Objects/abstract.c b/Objects/abstract.c
index 0105c5d..a68253b 100644
--- a/Objects/abstract.c
+++ b/Objects/abstract.c
@@ -169,11 +169,11 @@ PyObject_GetItem(PyObject *o, PyObject *key)
     }
 
     if (PyType_Check(o)) {
-        PyObject *meth, *result, *stack[2] = {o, key};
+        PyObject *meth, *result, *stack[1] = {key};
         _Py_IDENTIFIER(__class_getitem__);
         meth = _PyObject_GetAttrId(o, &PyId___class_getitem__);
         if (meth) {
-            result = _PyObject_FastCall(meth, stack, 2);
+            result = _PyObject_FastCall(meth, stack, 1);
             Py_DECREF(meth);
             return result;
         }
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index e570c41..0c8f0ad 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -55,6 +55,7 @@ static size_t method_cache_collisions = 0;
 /* alphabetical order */
 _Py_IDENTIFIER(__abstractmethods__);
 _Py_IDENTIFIER(__class__);
+_Py_IDENTIFIER(__class_getitem__);
 _Py_IDENTIFIER(__delitem__);
 _Py_IDENTIFIER(__dict__);
 _Py_IDENTIFIER(__doc__);
@@ -2694,8 +2695,8 @@ type_new(PyTypeObject *metatype, PyObject *args, PyObject *kwds)
         Py_DECREF(tmp);
     }
 
-    /* Special-case __init_subclass__: if it's a plain function,
-       make it a classmethod */
+    /* Special-case __init_subclass__ and __class_getitem__:
+       if they are plain functions, make them classmethods */
     tmp = _PyDict_GetItemId(dict, &PyId___init_subclass__);
     if (tmp != NULL && PyFunction_Check(tmp)) {
         tmp = PyClassMethod_New(tmp);
@@ -2708,6 +2709,18 @@ type_new(PyTypeObject *metatype, PyObject *args, PyObject *kwds)
         Py_DECREF(tmp);
     }
 
+    tmp = _PyDict_GetItemId(dict, &PyId___class_getitem__);
+    if (tmp != NULL && PyFunction_Check(tmp)) {
+        tmp = PyClassMethod_New(tmp);
+        if (tmp == NULL)
+            goto error;
+        if (_PyDict_SetItemId(dict, &PyId___class_getitem__, tmp) < 0) {
+            Py_DECREF(tmp);
+            goto error;
+        }
+        Py_DECREF(tmp);
+    }
+
     /* Add descriptors for custom slots from __slots__, or for __dict__ */
     mp = PyHeapType_GET_MEMBERS(et);
     slotoffset = base->tp_basicsize;