Skip to content

Commit f8e7d88

Browse files
committed
Add vectorcall for itemgetter objects
1 parent 3240bc6 commit f8e7d88

File tree

1 file changed

+38
-5
lines changed

1 file changed

+38
-5
lines changed

Modules/_operator.c

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "Python.h"
22
#include "pycore_moduleobject.h" // _PyModule_GetState()
3+
#include "structmember.h" // PyMemberDef
34
#include "clinic/_operator.c.h"
45

56
typedef struct {
@@ -953,8 +954,15 @@ typedef struct {
953954
Py_ssize_t nitems;
954955
PyObject *item;
955956
Py_ssize_t index; // -1 unless *item* is a single non-negative integer index
957+
vectorcallfunc vectorcall;
956958
} itemgetterobject;
957959

960+
// Forward declarations
961+
static PyObject *
962+
itemgetter_vectorcall(PyObject *, PyObject *const *, size_t, PyObject *);
963+
static PyObject *
964+
itemgetter_call_impl(itemgetterobject *, PyObject *);
965+
958966
/* AC 3.5: treats first argument as an iterable, otherwise uses *args */
959967
static PyObject *
960968
itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
@@ -1000,6 +1008,7 @@ itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
10001008
}
10011009
}
10021010

1011+
ig->vectorcall = (vectorcallfunc)itemgetter_vectorcall;
10031012
PyObject_GC_Track(ig);
10041013
return (PyObject *)ig;
10051014
}
@@ -1032,16 +1041,33 @@ itemgetter_traverse(itemgetterobject *ig, visitproc visit, void *arg)
10321041
static PyObject *
10331042
itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw)
10341043
{
1035-
PyObject *obj, *result;
1036-
Py_ssize_t i, nitems=ig->nitems;
1037-
10381044
assert(PyTuple_CheckExact(args));
10391045
if (!_PyArg_NoKeywords("itemgetter", kw))
10401046
return NULL;
10411047
if (!_PyArg_CheckPositional("itemgetter", PyTuple_GET_SIZE(args), 1, 1))
10421048
return NULL;
1049+
return itemgetter_call_impl(ig, PyTuple_GET_ITEM(args, 0));
1050+
}
10431051

1044-
obj = PyTuple_GET_ITEM(args, 0);
1052+
static PyObject *
1053+
itemgetter_vectorcall(PyObject *ig, PyObject *const *args,
1054+
size_t nargsf, PyObject *kwnames)
1055+
{
1056+
if (!_PyArg_NoKwnames("itemgetter", kwnames)) {
1057+
return NULL;
1058+
}
1059+
Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
1060+
if (!_PyArg_CheckPositional("itemgetter", nargs, 1, 1)) {
1061+
return NULL;
1062+
}
1063+
return itemgetter_call_impl((itemgetterobject *)ig, args[0]);
1064+
}
1065+
1066+
static PyObject *
1067+
itemgetter_call_impl(itemgetterobject *ig, PyObject *obj)
1068+
{
1069+
PyObject *result;
1070+
Py_ssize_t i, nitems=ig->nitems;
10451071
if (nitems == 1) {
10461072
if (ig->index >= 0
10471073
&& PyTuple_CheckExact(obj)
@@ -1109,20 +1135,27 @@ static PyMethodDef itemgetter_methods[] = {
11091135
{NULL}
11101136
};
11111137

1138+
static PyMemberDef itemgetter_members[] = {
1139+
{"__vectorcalloffset__", T_PYSSIZET, offsetof(itemgetterobject, vectorcall), READONLY},
1140+
{NULL} /* Sentinel */
1141+
};
1142+
11121143
PyDoc_STRVAR(itemgetter_doc,
11131144
"itemgetter(item, ...) --> itemgetter object\n\
11141145
\n\
11151146
Return a callable object that fetches the given item(s) from its operand.\n\
11161147
After f = itemgetter(2), the call f(r) returns r[2].\n\
11171148
After g = itemgetter(2, 5, 3), the call g(r) returns (r[2], r[5], r[3])");
11181149

1150+
11191151
static PyType_Slot itemgetter_type_slots[] = {
11201152
{Py_tp_doc, (void *)itemgetter_doc},
11211153
{Py_tp_dealloc, itemgetter_dealloc},
11221154
{Py_tp_call, itemgetter_call},
11231155
{Py_tp_traverse, itemgetter_traverse},
11241156
{Py_tp_clear, itemgetter_clear},
11251157
{Py_tp_methods, itemgetter_methods},
1158+
{Py_tp_members, itemgetter_members},
11261159
{Py_tp_new, itemgetter_new},
11271160
{Py_tp_getattro, PyObject_GenericGetAttr},
11281161
{Py_tp_repr, itemgetter_repr},
@@ -1134,7 +1167,7 @@ static PyType_Spec itemgetter_type_spec = {
11341167
.basicsize = sizeof(itemgetterobject),
11351168
.itemsize = 0,
11361169
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1137-
Py_TPFLAGS_IMMUTABLETYPE),
1170+
Py_TPFLAGS_IMMUTABLETYPE | Py_TPFLAGS_HAVE_VECTORCALL),
11381171
.slots = itemgetter_type_slots,
11391172
};
11401173

0 commit comments

Comments
 (0)