diff --git a/Include/cpython/object.h b/Include/cpython/object.h index ae7f780a931..8f9153c59a1 100644 --- a/Include/cpython/object.h +++ b/Include/cpython/object.h @@ -132,6 +132,18 @@ typedef struct { sendfunc am_send; } PyAsyncMethods; +typedef void (*setawaiterfunc)(PyObject *receiver, PyObject *awaiter); + +typedef struct { + PyAsyncMethods ame_async_methods; + setawaiterfunc ame_setawaiter; +} Ci_AsyncMethodsWithExtra; + +#define Ci_HeapType_AM_EXTRA(etype) \ + ((Ci_AsyncMethodsWithExtra *)(((char *)etype) + \ + Py_TYPE(etype)->tp_basicsize + \ + Py_SIZE(etype) * sizeof(PyMemberDef))) + typedef struct { getbufferproc bf_getbuffer; releasebufferproc bf_releasebuffer; diff --git a/Include/object.h b/Include/object.h index 0d94cf82553..8494114514f 100644 --- a/Include/object.h +++ b/Include/object.h @@ -527,6 +527,9 @@ given type object has a specified feature. /* Objects behave like an unbound method */ #define Py_TPFLAGS_METHOD_DESCRIPTOR (1UL << 17) +/* Set if type's tp_as_async slot points to Ci_AsyncMethodsWithExtra */ +#define Ci_TPFLAGS_HAVE_AM_EXTRA (1UL << 18) + /* Object has up-to-date type attribute cache */ #define Py_TPFLAGS_VALID_VERSION_TAG (1UL << 19) diff --git a/Objects/typeobject.c b/Objects/typeobject.c index 4b63b20bc2c..a353e85a85e 100644 --- a/Objects/typeobject.c +++ b/Objects/typeobject.c @@ -3318,21 +3318,45 @@ type_new_alloc(type_new_ctx *ctx) PyTypeObject *metatype = ctx->metatype; PyTypeObject *type; + /* If the base class has Ci_AsyncMethodsWithExtra, we allocate space at the + * end of this type so it can also have it. The extra slots in + * Ci_AsyncMethodsWithExtra aren't added to slotdefs and don't have managed + * counterparts, so this has no implications on slotptr() or the relative + * order of the various *Methods members of PyHeapTypeObject. */ + int have_am_extra = PyType_HasFeature(ctx->base, Ci_TPFLAGS_HAVE_AM_EXTRA); + + /* Allocate the type object */ + Py_ssize_t extra_bytes = + (have_am_extra ? sizeof(Ci_AsyncMethodsWithExtra) : 0); + Py_ssize_t extra_slots = + (extra_bytes + sizeof(PyMemberDef) - 1) / sizeof(PyMemberDef); + // Allocate the type object - type = (PyTypeObject *)metatype->tp_alloc(metatype, ctx->nslot); + type = (PyTypeObject *)metatype->tp_alloc(metatype, ctx->nslot + extra_slots); if (type == NULL) { return NULL; } + Py_SET_SIZE(type, Py_SIZE(type) - extra_slots); + PyHeapTypeObject *et = (PyHeapTypeObject *)type; // Initialize tp_flags. // All heap types need GC, since we can create a reference cycle by storing // an instance on one of its parents. type->tp_flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HEAPTYPE | - Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC); + Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC | + (ctx->base->tp_flags & Ci_TPFLAGS_HAVE_AM_EXTRA)); // Initialize essential fields - type->tp_as_async = &et->as_async; + if (have_am_extra) { + type->tp_as_async = (PyAsyncMethods *)Ci_HeapType_AM_EXTRA(type); + /* Only ame_setawaiter is inherited and it has no managed counterpart, + * so it's special-cased here. */ + ((Ci_AsyncMethodsWithExtra *)type->tp_as_async)->ame_setawaiter = + ((Ci_AsyncMethodsWithExtra *)ctx->base->tp_as_async)->ame_setawaiter; + } else { + type->tp_as_async = &et->as_async; + } type->tp_as_number = &et->as_number; type->tp_as_sequence = &et->as_sequence; type->tp_as_mapping = &et->as_mapping;