bpo-36974: inherit the vectorcall protocol (GH-13498)

3 files changed, 104 insertions, 2 deletions

diff --git a/Lib/test/test_capi.py b/Lib/test/test_capi.py
index 0813abb..795aa78 100644
--- a/Lib/test/test_capi.py
+++ b/Lib/test/test_capi.py
@@ -27,6 +27,7 @@ _testcapi = support.import_module('_testcapi')
 # Were we compiled --with-pydebug or with #define Py_DEBUG?
 Py_DEBUG = hasattr(sys, 'gettotalrefcount')
 
+Py_TPFLAGS_HAVE_VECTORCALL = 1 << 11
 Py_TPFLAGS_METHOD_DESCRIPTOR = 1 << 17
 
 
@@ -484,6 +485,27 @@ class TestPEP590(unittest.TestCase):
             pass
         self.assertFalse(MethodDescriptorHeap.__flags__ & Py_TPFLAGS_METHOD_DESCRIPTOR)
 
+    def test_vectorcall_flag(self):
+        self.assertTrue(_testcapi.MethodDescriptorBase.__flags__ & Py_TPFLAGS_HAVE_VECTORCALL)
+        self.assertTrue(_testcapi.MethodDescriptorDerived.__flags__ & Py_TPFLAGS_HAVE_VECTORCALL)
+        self.assertFalse(_testcapi.MethodDescriptorNopGet.__flags__ & Py_TPFLAGS_HAVE_VECTORCALL)
+        self.assertTrue(_testcapi.MethodDescriptor2.__flags__ & Py_TPFLAGS_HAVE_VECTORCALL)
+
+        # Heap type should not inherit Py_TPFLAGS_HAVE_VECTORCALL
+        class MethodDescriptorHeap(_testcapi.MethodDescriptorBase):
+            pass
+        self.assertFalse(MethodDescriptorHeap.__flags__ & Py_TPFLAGS_HAVE_VECTORCALL)
+
+    def test_vectorcall_override(self):
+        # Check that tp_call can correctly override vectorcall.
+        # MethodDescriptorNopGet implements tp_call but it inherits from
+        # MethodDescriptorBase, which implements vectorcall. Since
+        # MethodDescriptorNopGet returns the args tuple when called, we check
+        # additionally that no new tuple is created for this call.
+        args = tuple(range(5))
+        f = _testcapi.MethodDescriptorNopGet()
+        self.assertIs(f(*args), args)
+
     def test_vectorcall(self):
         # Test a bunch of different ways to call objects:
         # 1. normal call
@@ -498,7 +520,10 @@ class TestPEP590(unittest.TestCase):
                  ([].append, (0,), {}, None),
                  (sum, ([36],), {"start":6}, 42),
                  (testfunction, (42,), {}, 42),
-                 (testfunction_kw, (42,), {"kw":None}, 42)]
+                 (testfunction_kw, (42,), {"kw":None}, 42),
+                 (_testcapi.MethodDescriptorBase(), (0,), {}, True),
+                 (_testcapi.MethodDescriptorDerived(), (0,), {}, True),
+                 (_testcapi.MethodDescriptor2(), (0,), {}, False)]
 
         from _testcapi import pyobject_vectorcall, pyvectorcall_call
         from types import MethodType
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index f2f418c..a7451c6 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -5814,6 +5814,29 @@ static PyTypeObject Generic_Type = {
 
 /* Test PEP 590 */
 
+typedef struct {
+    PyObject_HEAD
+    vectorcallfunc vectorcall;
+} MethodDescriptorObject;
+
+static PyObject *
+MethodDescriptor_vectorcall(PyObject *callable, PyObject *const *args,
+                            size_t nargsf, PyObject *kwnames)
+{
+    /* True if using the vectorcall function in MethodDescriptorObject
+     * but False for MethodDescriptor2Object */
+    MethodDescriptorObject *md = (MethodDescriptorObject *)callable;
+    return PyBool_FromLong(md->vectorcall != NULL);
+}
+
+static PyObject *
+MethodDescriptor_new(PyTypeObject* type, PyObject* args, PyObject *kw)
+{
+    MethodDescriptorObject *op = PyObject_New(MethodDescriptorObject, type);
+    op->vectorcall = MethodDescriptor_vectorcall;
+    return (PyObject *)op;
+}
+
 static PyObject *
 func_descr_get(PyObject *func, PyObject *obj, PyObject *type)
 {
@@ -5831,10 +5854,22 @@ nop_descr_get(PyObject *func, PyObject *obj, PyObject *type)
     return func;
 }
 
+static PyObject *
+call_return_args(PyObject *self, PyObject *args, PyObject *kwargs)
+{
+    Py_INCREF(args);
+    return args;
+}
+
 static PyTypeObject MethodDescriptorBase_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     "MethodDescriptorBase",
-    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_METHOD_DESCRIPTOR,
+    sizeof(MethodDescriptorObject),
+    .tp_new = MethodDescriptor_new,
+    .tp_call = PyVectorcall_Call,
+    .tp_vectorcall_offset = offsetof(MethodDescriptorObject, vectorcall),
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE |
+                Py_TPFLAGS_METHOD_DESCRIPTOR | _Py_TPFLAGS_HAVE_VECTORCALL,
     .tp_descr_get = func_descr_get,
 };
 
@@ -5848,9 +5883,34 @@ static PyTypeObject MethodDescriptorNopGet_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     "MethodDescriptorNopGet",
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+    .tp_call = call_return_args,
     .tp_descr_get = nop_descr_get,
 };
 
+typedef struct {
+    MethodDescriptorObject base;
+    vectorcallfunc vectorcall;
+} MethodDescriptor2Object;
+
+static PyObject *
+MethodDescriptor2_new(PyTypeObject* type, PyObject* args, PyObject *kw)
+{
+    MethodDescriptor2Object *op = PyObject_New(MethodDescriptor2Object, type);
+    op->base.vectorcall = NULL;
+    op->vectorcall = MethodDescriptor_vectorcall;
+    return (PyObject *)op;
+}
+
+static PyTypeObject MethodDescriptor2_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "MethodDescriptor2",
+    sizeof(MethodDescriptor2Object),
+    .tp_new = MethodDescriptor2_new,
+    .tp_call = PyVectorcall_Call,
+    .tp_vectorcall_offset = offsetof(MethodDescriptor2Object, vectorcall),
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | _Py_TPFLAGS_HAVE_VECTORCALL,
+};
+
 
 static struct PyModuleDef _testcapimodule = {
     PyModuleDef_HEAD_INIT,
@@ -5916,6 +5976,12 @@ PyInit__testcapi(void)
     Py_INCREF(&MethodDescriptorNopGet_Type);
     PyModule_AddObject(m, "MethodDescriptorNopGet", (PyObject *)&MethodDescriptorNopGet_Type);
 
+    MethodDescriptor2_Type.tp_base = &MethodDescriptorBase_Type;
+    if (PyType_Ready(&MethodDescriptor2_Type) < 0)
+        return NULL;
+    Py_INCREF(&MethodDescriptor2_Type);
+    PyModule_AddObject(m, "MethodDescriptor2", (PyObject *)&MethodDescriptor2_Type);
+
     if (PyType_Ready(&GenericAlias_Type) < 0)
         return NULL;
     Py_INCREF(&GenericAlias_Type);
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 071ff27..ac5a686 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -5147,6 +5147,17 @@ inherit_slots(PyTypeObject *type, PyTypeObject *base)
     COPYSLOT(tp_repr);
     /* tp_hash see tp_richcompare */
     COPYSLOT(tp_call);
+    /* Inherit tp_vectorcall_offset and _Py_TPFLAGS_HAVE_VECTORCALL if tp_call
+     * was inherited, but only for extension types */
+    if ((base->tp_flags & _Py_TPFLAGS_HAVE_VECTORCALL) &&
+        !(type->tp_flags & _Py_TPFLAGS_HAVE_VECTORCALL) &&
+        !(type->tp_flags & Py_TPFLAGS_HEAPTYPE) &&
+        base->tp_call &&
+        type->tp_call == base->tp_call)
+    {
+        type->tp_vectorcall_offset = base->tp_vectorcall_offset;
+        type->tp_flags |= _Py_TPFLAGS_HAVE_VECTORCALL;
+    }
     COPYSLOT(tp_str);
     {
         /* Copy comparison-related slots only when