bpo-43706: Use PEP 590 vectorcall to speed up enumerate() (GH-25154)
diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-04-03-02-44-15.bpo-43706.jjsXlT.rst b/Misc/NEWS.d/next/Core and Builtins/2021-04-03-02-44-15.bpo-43706.jjsXlT.rst
new file mode 100644
index 0000000..ee859e3
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2021-04-03-02-44-15.bpo-43706.jjsXlT.rst
@@ -0,0 +1,2 @@
+Speed up calls to ``enumerate()`` by using the :pep:`590` ``vectorcall``
+calling convention. Patch by Dong-hee Na.
diff --git a/Objects/enumobject.c b/Objects/enumobject.c
index 4513831..b78230d 100644
--- a/Objects/enumobject.c
+++ b/Objects/enumobject.c
@@ -81,6 +81,45 @@ enum_new_impl(PyTypeObject *type, PyObject *iterable, PyObject *start)
return (PyObject *)en;
}
+// TODO: Use AC when bpo-43447 is supported
+static PyObject *
+enumerate_vectorcall(PyObject *type, PyObject *const *args,
+ size_t nargsf, PyObject *kwnames)
+{
+ assert(PyType_Check(type));
+ PyTypeObject *tp = (PyTypeObject *)type;
+ Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
+ Py_ssize_t nkwargs = 0;
+ if (nargs == 0) {
+ PyErr_SetString(PyExc_TypeError,
+ "enumerate() missing required argument 'iterable'");
+ return NULL;
+ }
+ if (kwnames != NULL) {
+ nkwargs = PyTuple_GET_SIZE(kwnames);
+ }
+
+ if (nargs + nkwargs == 2) {
+ if (nkwargs == 1) {
+ PyObject *kw = PyTuple_GET_ITEM(kwnames, 0);
+ if (!_PyUnicode_EqualToASCIIString(kw, "start")) {
+ PyErr_Format(PyExc_TypeError,
+ "'%S' is an invalid keyword argument for enumerate()", kw);
+ return NULL;
+ }
+ }
+ return enum_new_impl(tp, args[0], args[1]);
+ }
+
+ if (nargs == 1 && nkwargs == 0) {
+ return enum_new_impl(tp, args[0], NULL);
+ }
+
+ PyErr_Format(PyExc_TypeError,
+ "enumerate() takes at most 2 arguments (%d given)", nargs + nkwargs);
+ return NULL;
+}
+
static void
enum_dealloc(enumobject *en)
{
@@ -261,6 +300,7 @@ PyTypeObject PyEnum_Type = {
PyType_GenericAlloc, /* tp_alloc */
enum_new, /* tp_new */
PyObject_GC_Del, /* tp_free */
+ .tp_vectorcall = (vectorcallfunc)enumerate_vectorcall
};
/* Reversed Object ***************************************************************/
