1
2 #include "Python.h"
3
4 PyDoc_STRVAR(operator_doc,
5 "Operator interface.\n\
6 \n\
7 This module exports a set of functions implemented in C corresponding\n\
8 to the intrinsic operators of Python. For example, operator.add(x, y)\n\
9 is equivalent to the expression x+y. The function names are those\n\
10 used for special methods; variants without leading and trailing\n\
11 '__' are also provided for convenience.");
12
13 #define spam1(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a1) { \
14 return AOP(a1); }
15
16 #define spam2(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
17 PyObject *a1, *a2; \
18 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
19 return AOP(a1,a2); }
20
21 #define spamoi(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
22 PyObject *a1; int a2; \
23 if(! PyArg_ParseTuple(a,"Oi:" #OP,&a1,&a2)) return NULL; \
24 return AOP(a1,a2); }
25
26 #define spam2n(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
27 PyObject *a1, *a2; \
28 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
29 if(-1 == AOP(a1,a2)) return NULL; \
30 Py_INCREF(Py_None); \
31 return Py_None; }
32
33 #define spam3n(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
34 PyObject *a1, *a2, *a3; \
35 if(! PyArg_UnpackTuple(a,#OP,3,3,&a1,&a2,&a3)) return NULL; \
36 if(-1 == AOP(a1,a2,a3)) return NULL; \
37 Py_INCREF(Py_None); \
38 return Py_None; }
39
40 #define spami(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a1) { \
41 long r; \
42 if(-1 == (r=AOP(a1))) return NULL; \
43 return PyBool_FromLong(r); }
44
45 #define spami2(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
46 PyObject *a1, *a2; long r; \
47 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
48 if(-1 == (r=AOP(a1,a2))) return NULL; \
49 return PyLong_FromLong(r); }
50
51 #define spamn2(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
52 PyObject *a1, *a2; Py_ssize_t r; \
53 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
54 if(-1 == (r=AOP(a1,a2))) return NULL; \
55 return PyLong_FromSsize_t(r); }
56
57 #define spami2b(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
58 PyObject *a1, *a2; long r; \
59 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
60 if(-1 == (r=AOP(a1,a2))) return NULL; \
61 return PyBool_FromLong(r); }
62
63 #define spamrc(OP,A) static PyObject *OP(PyObject *s, PyObject *a) { \
64 PyObject *a1, *a2; \
65 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
66 return PyObject_RichCompare(a1,a2,A); }
67
spami(truth,PyObject_IsTrue)68 spami(truth , PyObject_IsTrue)
69 spam2(op_add , PyNumber_Add)
70 spam2(op_sub , PyNumber_Subtract)
71 spam2(op_mul , PyNumber_Multiply)
72 spam2(op_matmul , PyNumber_MatrixMultiply)
73 spam2(op_floordiv , PyNumber_FloorDivide)
74 spam2(op_truediv , PyNumber_TrueDivide)
75 spam2(op_mod , PyNumber_Remainder)
76 spam1(op_neg , PyNumber_Negative)
77 spam1(op_pos , PyNumber_Positive)
78 spam1(op_abs , PyNumber_Absolute)
79 spam1(op_inv , PyNumber_Invert)
80 spam1(op_invert , PyNumber_Invert)
81 spam2(op_lshift , PyNumber_Lshift)
82 spam2(op_rshift , PyNumber_Rshift)
83 spami(op_not_ , PyObject_Not)
84 spam2(op_and_ , PyNumber_And)
85 spam2(op_xor , PyNumber_Xor)
86 spam2(op_or_ , PyNumber_Or)
87 spam2(op_iadd , PyNumber_InPlaceAdd)
88 spam2(op_isub , PyNumber_InPlaceSubtract)
89 spam2(op_imul , PyNumber_InPlaceMultiply)
90 spam2(op_imatmul , PyNumber_InPlaceMatrixMultiply)
91 spam2(op_ifloordiv , PyNumber_InPlaceFloorDivide)
92 spam2(op_itruediv , PyNumber_InPlaceTrueDivide)
93 spam2(op_imod , PyNumber_InPlaceRemainder)
94 spam2(op_ilshift , PyNumber_InPlaceLshift)
95 spam2(op_irshift , PyNumber_InPlaceRshift)
96 spam2(op_iand , PyNumber_InPlaceAnd)
97 spam2(op_ixor , PyNumber_InPlaceXor)
98 spam2(op_ior , PyNumber_InPlaceOr)
99 spam2(op_concat , PySequence_Concat)
100 spam2(op_iconcat , PySequence_InPlaceConcat)
101 spami2b(op_contains , PySequence_Contains)
102 spamn2(indexOf , PySequence_Index)
103 spamn2(countOf , PySequence_Count)
104 spam2(op_getitem , PyObject_GetItem)
105 spam2n(op_delitem , PyObject_DelItem)
106 spam3n(op_setitem , PyObject_SetItem)
107 spamrc(op_lt , Py_LT)
108 spamrc(op_le , Py_LE)
109 spamrc(op_eq , Py_EQ)
110 spamrc(op_ne , Py_NE)
111 spamrc(op_gt , Py_GT)
112 spamrc(op_ge , Py_GE)
113
114 static PyObject*
115 op_pow(PyObject *s, PyObject *a)
116 {
117 PyObject *a1, *a2;
118 if (PyArg_UnpackTuple(a,"pow", 2, 2, &a1, &a2))
119 return PyNumber_Power(a1, a2, Py_None);
120 return NULL;
121 }
122
123 static PyObject*
op_ipow(PyObject * s,PyObject * a)124 op_ipow(PyObject *s, PyObject *a)
125 {
126 PyObject *a1, *a2;
127 if (PyArg_UnpackTuple(a,"ipow", 2, 2, &a1, &a2))
128 return PyNumber_InPlacePower(a1, a2, Py_None);
129 return NULL;
130 }
131
132 static PyObject *
op_index(PyObject * s,PyObject * a)133 op_index(PyObject *s, PyObject *a)
134 {
135 return PyNumber_Index(a);
136 }
137
138 static PyObject*
is_(PyObject * s,PyObject * a)139 is_(PyObject *s, PyObject *a)
140 {
141 PyObject *a1, *a2, *result = NULL;
142 if (PyArg_UnpackTuple(a,"is_", 2, 2, &a1, &a2)) {
143 result = (a1 == a2) ? Py_True : Py_False;
144 Py_INCREF(result);
145 }
146 return result;
147 }
148
149 static PyObject*
is_not(PyObject * s,PyObject * a)150 is_not(PyObject *s, PyObject *a)
151 {
152 PyObject *a1, *a2, *result = NULL;
153 if (PyArg_UnpackTuple(a,"is_not", 2, 2, &a1, &a2)) {
154 result = (a1 != a2) ? Py_True : Py_False;
155 Py_INCREF(result);
156 }
157 return result;
158 }
159
160 #undef spam1
161 #undef spam2
162 #undef spam1o
163 #undef spam1o
164
165 /* compare_digest **********************************************************/
166
167 /*
168 * timing safe compare
169 *
170 * Returns 1 of the strings are equal.
171 * In case of len(a) != len(b) the function tries to keep the timing
172 * dependent on the length of b. CPU cache locally may still alter timing
173 * a bit.
174 */
175 static int
_tscmp(const unsigned char * a,const unsigned char * b,Py_ssize_t len_a,Py_ssize_t len_b)176 _tscmp(const unsigned char *a, const unsigned char *b,
177 Py_ssize_t len_a, Py_ssize_t len_b)
178 {
179 /* The volatile type declarations make sure that the compiler has no
180 * chance to optimize and fold the code in any way that may change
181 * the timing.
182 */
183 volatile Py_ssize_t length;
184 volatile const unsigned char *left;
185 volatile const unsigned char *right;
186 Py_ssize_t i;
187 unsigned char result;
188
189 /* loop count depends on length of b */
190 length = len_b;
191 left = NULL;
192 right = b;
193
194 /* don't use else here to keep the amount of CPU instructions constant,
195 * volatile forces re-evaluation
196 * */
197 if (len_a == length) {
198 left = *((volatile const unsigned char**)&a);
199 result = 0;
200 }
201 if (len_a != length) {
202 left = b;
203 result = 1;
204 }
205
206 for (i=0; i < length; i++) {
207 result |= *left++ ^ *right++;
208 }
209
210 return (result == 0);
211 }
212
213 PyDoc_STRVAR(length_hint__doc__,
214 "length_hint(obj, default=0) -> int\n"
215 "Return an estimate of the number of items in obj.\n"
216 "This is useful for presizing containers when building from an\n"
217 "iterable.\n"
218 "\n"
219 "If the object supports len(), the result will be\n"
220 "exact. Otherwise, it may over- or under-estimate by an\n"
221 "arbitrary amount. The result will be an integer >= 0.");
222
length_hint(PyObject * self,PyObject * args)223 static PyObject *length_hint(PyObject *self, PyObject *args)
224 {
225 PyObject *obj;
226 Py_ssize_t defaultvalue = 0, res;
227 if (!PyArg_ParseTuple(args, "O|n:length_hint", &obj, &defaultvalue)) {
228 return NULL;
229 }
230 res = PyObject_LengthHint(obj, defaultvalue);
231 if (res == -1 && PyErr_Occurred()) {
232 return NULL;
233 }
234 return PyLong_FromSsize_t(res);
235 }
236
237
238 PyDoc_STRVAR(compare_digest__doc__,
239 "compare_digest(a, b) -> bool\n"
240 "\n"
241 "Return 'a == b'. This function uses an approach designed to prevent\n"
242 "timing analysis, making it appropriate for cryptography.\n"
243 "a and b must both be of the same type: either str (ASCII only),\n"
244 "or any bytes-like object.\n"
245 "\n"
246 "Note: If a and b are of different lengths, or if an error occurs,\n"
247 "a timing attack could theoretically reveal information about the\n"
248 "types and lengths of a and b--but not their values.\n");
249
250 static PyObject*
compare_digest(PyObject * self,PyObject * args)251 compare_digest(PyObject *self, PyObject *args)
252 {
253 PyObject *a, *b;
254 int rc;
255
256 if (!PyArg_ParseTuple(args, "OO:compare_digest", &a, &b)) {
257 return NULL;
258 }
259
260 /* ASCII unicode string */
261 if(PyUnicode_Check(a) && PyUnicode_Check(b)) {
262 if (PyUnicode_READY(a) == -1 || PyUnicode_READY(b) == -1) {
263 return NULL;
264 }
265 if (!PyUnicode_IS_ASCII(a) || !PyUnicode_IS_ASCII(b)) {
266 PyErr_SetString(PyExc_TypeError,
267 "comparing strings with non-ASCII characters is "
268 "not supported");
269 return NULL;
270 }
271
272 rc = _tscmp(PyUnicode_DATA(a),
273 PyUnicode_DATA(b),
274 PyUnicode_GET_LENGTH(a),
275 PyUnicode_GET_LENGTH(b));
276 }
277 /* fallback to buffer interface for bytes, bytesarray and other */
278 else {
279 Py_buffer view_a;
280 Py_buffer view_b;
281
282 if (PyObject_CheckBuffer(a) == 0 && PyObject_CheckBuffer(b) == 0) {
283 PyErr_Format(PyExc_TypeError,
284 "unsupported operand types(s) or combination of types: "
285 "'%.100s' and '%.100s'",
286 Py_TYPE(a)->tp_name, Py_TYPE(b)->tp_name);
287 return NULL;
288 }
289
290 if (PyObject_GetBuffer(a, &view_a, PyBUF_SIMPLE) == -1) {
291 return NULL;
292 }
293 if (view_a.ndim > 1) {
294 PyErr_SetString(PyExc_BufferError,
295 "Buffer must be single dimension");
296 PyBuffer_Release(&view_a);
297 return NULL;
298 }
299
300 if (PyObject_GetBuffer(b, &view_b, PyBUF_SIMPLE) == -1) {
301 PyBuffer_Release(&view_a);
302 return NULL;
303 }
304 if (view_b.ndim > 1) {
305 PyErr_SetString(PyExc_BufferError,
306 "Buffer must be single dimension");
307 PyBuffer_Release(&view_a);
308 PyBuffer_Release(&view_b);
309 return NULL;
310 }
311
312 rc = _tscmp((const unsigned char*)view_a.buf,
313 (const unsigned char*)view_b.buf,
314 view_a.len,
315 view_b.len);
316
317 PyBuffer_Release(&view_a);
318 PyBuffer_Release(&view_b);
319 }
320
321 return PyBool_FromLong(rc);
322 }
323
324 /* operator methods **********************************************************/
325
326 #define spam1(OP,DOC) {#OP, OP, METH_VARARGS, PyDoc_STR(DOC)},
327 #define spam2(OP,DOC) {#OP, op_##OP, METH_VARARGS, PyDoc_STR(DOC)},
328 #define spam1o(OP,DOC) {#OP, OP, METH_O, PyDoc_STR(DOC)},
329 #define spam2o(OP,DOC) {#OP, op_##OP, METH_O, PyDoc_STR(DOC)},
330
331 static struct PyMethodDef operator_methods[] = {
332
333 spam1o(truth,
334 "truth(a) -- Return True if a is true, False otherwise.")
335 spam2(contains,
336 "contains(a, b) -- Same as b in a (note reversed operands).")
337 spam1(indexOf,
338 "indexOf(a, b) -- Return the first index of b in a.")
339 spam1(countOf,
340 "countOf(a, b) -- Return the number of times b occurs in a.")
341
342 spam1(is_, "is_(a, b) -- Same as a is b.")
343 spam1(is_not, "is_not(a, b) -- Same as a is not b.")
344 spam2o(index, "index(a) -- Same as a.__index__()")
345 spam2(add, "add(a, b) -- Same as a + b.")
346 spam2(sub, "sub(a, b) -- Same as a - b.")
347 spam2(mul, "mul(a, b) -- Same as a * b.")
348 spam2(matmul, "matmul(a, b) -- Same as a @ b.")
349 spam2(floordiv, "floordiv(a, b) -- Same as a // b.")
350 spam2(truediv, "truediv(a, b) -- Same as a / b.")
351 spam2(mod, "mod(a, b) -- Same as a % b.")
352 spam2o(neg, "neg(a) -- Same as -a.")
353 spam2o(pos, "pos(a) -- Same as +a.")
354 spam2o(abs, "abs(a) -- Same as abs(a).")
355 spam2o(inv, "inv(a) -- Same as ~a.")
356 spam2o(invert, "invert(a) -- Same as ~a.")
357 spam2(lshift, "lshift(a, b) -- Same as a << b.")
358 spam2(rshift, "rshift(a, b) -- Same as a >> b.")
359 spam2o(not_, "not_(a) -- Same as not a.")
360 spam2(and_, "and_(a, b) -- Same as a & b.")
361 spam2(xor, "xor(a, b) -- Same as a ^ b.")
362 spam2(or_, "or_(a, b) -- Same as a | b.")
363 spam2(iadd, "a = iadd(a, b) -- Same as a += b.")
364 spam2(isub, "a = isub(a, b) -- Same as a -= b.")
365 spam2(imul, "a = imul(a, b) -- Same as a *= b.")
366 spam2(imatmul, "a = imatmul(a, b) -- Same as a @= b.")
367 spam2(ifloordiv, "a = ifloordiv(a, b) -- Same as a //= b.")
368 spam2(itruediv, "a = itruediv(a, b) -- Same as a /= b")
369 spam2(imod, "a = imod(a, b) -- Same as a %= b.")
370 spam2(ilshift, "a = ilshift(a, b) -- Same as a <<= b.")
371 spam2(irshift, "a = irshift(a, b) -- Same as a >>= b.")
372 spam2(iand, "a = iand(a, b) -- Same as a &= b.")
373 spam2(ixor, "a = ixor(a, b) -- Same as a ^= b.")
374 spam2(ior, "a = ior(a, b) -- Same as a |= b.")
375 spam2(concat,
376 "concat(a, b) -- Same as a + b, for a and b sequences.")
377 spam2(iconcat,
378 "a = iconcat(a, b) -- Same as a += b, for a and b sequences.")
379 spam2(getitem,
380 "getitem(a, b) -- Same as a[b].")
381 spam2(setitem,
382 "setitem(a, b, c) -- Same as a[b] = c.")
383 spam2(delitem,
384 "delitem(a, b) -- Same as del a[b].")
385 spam2(pow, "pow(a, b) -- Same as a ** b.")
386 spam2(ipow, "a = ipow(a, b) -- Same as a **= b.")
387 spam2(lt, "lt(a, b) -- Same as a<b.")
388 spam2(le, "le(a, b) -- Same as a<=b.")
389 spam2(eq, "eq(a, b) -- Same as a==b.")
390 spam2(ne, "ne(a, b) -- Same as a!=b.")
391 spam2(gt, "gt(a, b) -- Same as a>b.")
392 spam2(ge, "ge(a, b) -- Same as a>=b.")
393
394 {"_compare_digest", (PyCFunction)compare_digest, METH_VARARGS,
395 compare_digest__doc__},
396 {"length_hint", (PyCFunction)length_hint, METH_VARARGS,
397 length_hint__doc__},
398 {NULL, NULL} /* sentinel */
399
400 };
401
402 /* itemgetter object **********************************************************/
403
404 typedef struct {
405 PyObject_HEAD
406 Py_ssize_t nitems;
407 PyObject *item;
408 } itemgetterobject;
409
410 static PyTypeObject itemgetter_type;
411
412 static PyObject *
itemgetter_new(PyTypeObject * type,PyObject * args,PyObject * kwds)413 itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
414 {
415 itemgetterobject *ig;
416 PyObject *item;
417 Py_ssize_t nitems;
418
419 if (!_PyArg_NoKeywords("itemgetter()", kwds))
420 return NULL;
421
422 nitems = PyTuple_GET_SIZE(args);
423 if (nitems <= 1) {
424 if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &item))
425 return NULL;
426 } else
427 item = args;
428
429 /* create itemgetterobject structure */
430 ig = PyObject_GC_New(itemgetterobject, &itemgetter_type);
431 if (ig == NULL)
432 return NULL;
433
434 Py_INCREF(item);
435 ig->item = item;
436 ig->nitems = nitems;
437
438 PyObject_GC_Track(ig);
439 return (PyObject *)ig;
440 }
441
442 static void
itemgetter_dealloc(itemgetterobject * ig)443 itemgetter_dealloc(itemgetterobject *ig)
444 {
445 PyObject_GC_UnTrack(ig);
446 Py_XDECREF(ig->item);
447 PyObject_GC_Del(ig);
448 }
449
450 static int
itemgetter_traverse(itemgetterobject * ig,visitproc visit,void * arg)451 itemgetter_traverse(itemgetterobject *ig, visitproc visit, void *arg)
452 {
453 Py_VISIT(ig->item);
454 return 0;
455 }
456
457 static PyObject *
itemgetter_call(itemgetterobject * ig,PyObject * args,PyObject * kw)458 itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw)
459 {
460 PyObject *obj, *result;
461 Py_ssize_t i, nitems=ig->nitems;
462
463 if (kw != NULL && !_PyArg_NoKeywords("itemgetter", kw))
464 return NULL;
465 if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj))
466 return NULL;
467 if (nitems == 1)
468 return PyObject_GetItem(obj, ig->item);
469
470 assert(PyTuple_Check(ig->item));
471 assert(PyTuple_GET_SIZE(ig->item) == nitems);
472
473 result = PyTuple_New(nitems);
474 if (result == NULL)
475 return NULL;
476
477 for (i=0 ; i < nitems ; i++) {
478 PyObject *item, *val;
479 item = PyTuple_GET_ITEM(ig->item, i);
480 val = PyObject_GetItem(obj, item);
481 if (val == NULL) {
482 Py_DECREF(result);
483 return NULL;
484 }
485 PyTuple_SET_ITEM(result, i, val);
486 }
487 return result;
488 }
489
490 static PyObject *
itemgetter_repr(itemgetterobject * ig)491 itemgetter_repr(itemgetterobject *ig)
492 {
493 PyObject *repr;
494 const char *reprfmt;
495
496 int status = Py_ReprEnter((PyObject *)ig);
497 if (status != 0) {
498 if (status < 0)
499 return NULL;
500 return PyUnicode_FromFormat("%s(...)", Py_TYPE(ig)->tp_name);
501 }
502
503 reprfmt = ig->nitems == 1 ? "%s(%R)" : "%s%R";
504 repr = PyUnicode_FromFormat(reprfmt, Py_TYPE(ig)->tp_name, ig->item);
505 Py_ReprLeave((PyObject *)ig);
506 return repr;
507 }
508
509 static PyObject *
itemgetter_reduce(itemgetterobject * ig)510 itemgetter_reduce(itemgetterobject *ig)
511 {
512 if (ig->nitems == 1)
513 return Py_BuildValue("O(O)", Py_TYPE(ig), ig->item);
514 return PyTuple_Pack(2, Py_TYPE(ig), ig->item);
515 }
516
517 PyDoc_STRVAR(reduce_doc, "Return state information for pickling");
518
519 static PyMethodDef itemgetter_methods[] = {
520 {"__reduce__", (PyCFunction)itemgetter_reduce, METH_NOARGS,
521 reduce_doc},
522 {NULL}
523 };
524
525 PyDoc_STRVAR(itemgetter_doc,
526 "itemgetter(item, ...) --> itemgetter object\n\
527 \n\
528 Return a callable object that fetches the given item(s) from its operand.\n\
529 After f = itemgetter(2), the call f(r) returns r[2].\n\
530 After g = itemgetter(2, 5, 3), the call g(r) returns (r[2], r[5], r[3])");
531
532 static PyTypeObject itemgetter_type = {
533 PyVarObject_HEAD_INIT(NULL, 0)
534 "operator.itemgetter", /* tp_name */
535 sizeof(itemgetterobject), /* tp_basicsize */
536 0, /* tp_itemsize */
537 /* methods */
538 (destructor)itemgetter_dealloc, /* tp_dealloc */
539 0, /* tp_print */
540 0, /* tp_getattr */
541 0, /* tp_setattr */
542 0, /* tp_reserved */
543 (reprfunc)itemgetter_repr, /* tp_repr */
544 0, /* tp_as_number */
545 0, /* tp_as_sequence */
546 0, /* tp_as_mapping */
547 0, /* tp_hash */
548 (ternaryfunc)itemgetter_call, /* tp_call */
549 0, /* tp_str */
550 PyObject_GenericGetAttr, /* tp_getattro */
551 0, /* tp_setattro */
552 0, /* tp_as_buffer */
553 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
554 itemgetter_doc, /* tp_doc */
555 (traverseproc)itemgetter_traverse, /* tp_traverse */
556 0, /* tp_clear */
557 0, /* tp_richcompare */
558 0, /* tp_weaklistoffset */
559 0, /* tp_iter */
560 0, /* tp_iternext */
561 itemgetter_methods, /* tp_methods */
562 0, /* tp_members */
563 0, /* tp_getset */
564 0, /* tp_base */
565 0, /* tp_dict */
566 0, /* tp_descr_get */
567 0, /* tp_descr_set */
568 0, /* tp_dictoffset */
569 0, /* tp_init */
570 0, /* tp_alloc */
571 itemgetter_new, /* tp_new */
572 0, /* tp_free */
573 };
574
575
576 /* attrgetter object **********************************************************/
577
578 typedef struct {
579 PyObject_HEAD
580 Py_ssize_t nattrs;
581 PyObject *attr;
582 } attrgetterobject;
583
584 static PyTypeObject attrgetter_type;
585
586 static PyObject *
attrgetter_new(PyTypeObject * type,PyObject * args,PyObject * kwds)587 attrgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
588 {
589 attrgetterobject *ag;
590 PyObject *attr;
591 Py_ssize_t nattrs, idx, char_idx;
592
593 if (!_PyArg_NoKeywords("attrgetter()", kwds))
594 return NULL;
595
596 nattrs = PyTuple_GET_SIZE(args);
597 if (nattrs <= 1) {
598 if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &attr))
599 return NULL;
600 }
601
602 attr = PyTuple_New(nattrs);
603 if (attr == NULL)
604 return NULL;
605
606 /* prepare attr while checking args */
607 for (idx = 0; idx < nattrs; ++idx) {
608 PyObject *item = PyTuple_GET_ITEM(args, idx);
609 Py_ssize_t item_len;
610 void *data;
611 unsigned int kind;
612 int dot_count;
613
614 if (!PyUnicode_Check(item)) {
615 PyErr_SetString(PyExc_TypeError,
616 "attribute name must be a string");
617 Py_DECREF(attr);
618 return NULL;
619 }
620 if (PyUnicode_READY(item)) {
621 Py_DECREF(attr);
622 return NULL;
623 }
624 item_len = PyUnicode_GET_LENGTH(item);
625 kind = PyUnicode_KIND(item);
626 data = PyUnicode_DATA(item);
627
628 /* check whethere the string is dotted */
629 dot_count = 0;
630 for (char_idx = 0; char_idx < item_len; ++char_idx) {
631 if (PyUnicode_READ(kind, data, char_idx) == '.')
632 ++dot_count;
633 }
634
635 if (dot_count == 0) {
636 Py_INCREF(item);
637 PyUnicode_InternInPlace(&item);
638 PyTuple_SET_ITEM(attr, idx, item);
639 } else { /* make it a tuple of non-dotted attrnames */
640 PyObject *attr_chain = PyTuple_New(dot_count + 1);
641 PyObject *attr_chain_item;
642 Py_ssize_t unibuff_from = 0;
643 Py_ssize_t unibuff_till = 0;
644 Py_ssize_t attr_chain_idx = 0;
645
646 if (attr_chain == NULL) {
647 Py_DECREF(attr);
648 return NULL;
649 }
650
651 for (; dot_count > 0; --dot_count) {
652 while (PyUnicode_READ(kind, data, unibuff_till) != '.') {
653 ++unibuff_till;
654 }
655 attr_chain_item = PyUnicode_Substring(item,
656 unibuff_from,
657 unibuff_till);
658 if (attr_chain_item == NULL) {
659 Py_DECREF(attr_chain);
660 Py_DECREF(attr);
661 return NULL;
662 }
663 PyUnicode_InternInPlace(&attr_chain_item);
664 PyTuple_SET_ITEM(attr_chain, attr_chain_idx, attr_chain_item);
665 ++attr_chain_idx;
666 unibuff_till = unibuff_from = unibuff_till + 1;
667 }
668
669 /* now add the last dotless name */
670 attr_chain_item = PyUnicode_Substring(item,
671 unibuff_from, item_len);
672 if (attr_chain_item == NULL) {
673 Py_DECREF(attr_chain);
674 Py_DECREF(attr);
675 return NULL;
676 }
677 PyUnicode_InternInPlace(&attr_chain_item);
678 PyTuple_SET_ITEM(attr_chain, attr_chain_idx, attr_chain_item);
679
680 PyTuple_SET_ITEM(attr, idx, attr_chain);
681 }
682 }
683
684 /* create attrgetterobject structure */
685 ag = PyObject_GC_New(attrgetterobject, &attrgetter_type);
686 if (ag == NULL) {
687 Py_DECREF(attr);
688 return NULL;
689 }
690
691 ag->attr = attr;
692 ag->nattrs = nattrs;
693
694 PyObject_GC_Track(ag);
695 return (PyObject *)ag;
696 }
697
698 static void
attrgetter_dealloc(attrgetterobject * ag)699 attrgetter_dealloc(attrgetterobject *ag)
700 {
701 PyObject_GC_UnTrack(ag);
702 Py_XDECREF(ag->attr);
703 PyObject_GC_Del(ag);
704 }
705
706 static int
attrgetter_traverse(attrgetterobject * ag,visitproc visit,void * arg)707 attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg)
708 {
709 Py_VISIT(ag->attr);
710 return 0;
711 }
712
713 static PyObject *
dotted_getattr(PyObject * obj,PyObject * attr)714 dotted_getattr(PyObject *obj, PyObject *attr)
715 {
716 PyObject *newobj;
717
718 /* attr is either a tuple or instance of str.
719 Ensured by the setup code of attrgetter_new */
720 if (PyTuple_CheckExact(attr)) { /* chained getattr */
721 Py_ssize_t name_idx = 0, name_count;
722 PyObject *attr_name;
723
724 name_count = PyTuple_GET_SIZE(attr);
725 Py_INCREF(obj);
726 for (name_idx = 0; name_idx < name_count; ++name_idx) {
727 attr_name = PyTuple_GET_ITEM(attr, name_idx);
728 newobj = PyObject_GetAttr(obj, attr_name);
729 Py_DECREF(obj);
730 if (newobj == NULL) {
731 return NULL;
732 }
733 /* here */
734 obj = newobj;
735 }
736 } else { /* single getattr */
737 newobj = PyObject_GetAttr(obj, attr);
738 if (newobj == NULL)
739 return NULL;
740 obj = newobj;
741 }
742
743 return obj;
744 }
745
746 static PyObject *
attrgetter_call(attrgetterobject * ag,PyObject * args,PyObject * kw)747 attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
748 {
749 PyObject *obj, *result;
750 Py_ssize_t i, nattrs=ag->nattrs;
751
752 if (kw != NULL && !_PyArg_NoKeywords("attrgetter", kw))
753 return NULL;
754 if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj))
755 return NULL;
756 if (ag->nattrs == 1) /* ag->attr is always a tuple */
757 return dotted_getattr(obj, PyTuple_GET_ITEM(ag->attr, 0));
758
759 assert(PyTuple_Check(ag->attr));
760 assert(PyTuple_GET_SIZE(ag->attr) == nattrs);
761
762 result = PyTuple_New(nattrs);
763 if (result == NULL)
764 return NULL;
765
766 for (i=0 ; i < nattrs ; i++) {
767 PyObject *attr, *val;
768 attr = PyTuple_GET_ITEM(ag->attr, i);
769 val = dotted_getattr(obj, attr);
770 if (val == NULL) {
771 Py_DECREF(result);
772 return NULL;
773 }
774 PyTuple_SET_ITEM(result, i, val);
775 }
776 return result;
777 }
778
779 static PyObject *
dotjoinattr(PyObject * attr,PyObject ** attrsep)780 dotjoinattr(PyObject *attr, PyObject **attrsep)
781 {
782 if (PyTuple_CheckExact(attr)) {
783 if (*attrsep == NULL) {
784 *attrsep = PyUnicode_FromString(".");
785 if (*attrsep == NULL)
786 return NULL;
787 }
788 return PyUnicode_Join(*attrsep, attr);
789 } else {
790 Py_INCREF(attr);
791 return attr;
792 }
793 }
794
795 static PyObject *
attrgetter_args(attrgetterobject * ag)796 attrgetter_args(attrgetterobject *ag)
797 {
798 Py_ssize_t i;
799 PyObject *attrsep = NULL;
800 PyObject *attrstrings = PyTuple_New(ag->nattrs);
801 if (attrstrings == NULL)
802 return NULL;
803
804 for (i = 0; i < ag->nattrs; ++i) {
805 PyObject *attr = PyTuple_GET_ITEM(ag->attr, i);
806 PyObject *attrstr = dotjoinattr(attr, &attrsep);
807 if (attrstr == NULL) {
808 Py_XDECREF(attrsep);
809 Py_DECREF(attrstrings);
810 return NULL;
811 }
812 PyTuple_SET_ITEM(attrstrings, i, attrstr);
813 }
814 Py_XDECREF(attrsep);
815 return attrstrings;
816 }
817
818 static PyObject *
attrgetter_repr(attrgetterobject * ag)819 attrgetter_repr(attrgetterobject *ag)
820 {
821 PyObject *repr = NULL;
822 int status = Py_ReprEnter((PyObject *)ag);
823 if (status != 0) {
824 if (status < 0)
825 return NULL;
826 return PyUnicode_FromFormat("%s(...)", Py_TYPE(ag)->tp_name);
827 }
828
829 if (ag->nattrs == 1) {
830 PyObject *attrsep = NULL;
831 PyObject *attr = dotjoinattr(PyTuple_GET_ITEM(ag->attr, 0), &attrsep);
832 if (attr != NULL) {
833 repr = PyUnicode_FromFormat("%s(%R)", Py_TYPE(ag)->tp_name, attr);
834 Py_DECREF(attr);
835 }
836 Py_XDECREF(attrsep);
837 }
838 else {
839 PyObject *attrstrings = attrgetter_args(ag);
840 if (attrstrings != NULL) {
841 repr = PyUnicode_FromFormat("%s%R",
842 Py_TYPE(ag)->tp_name, attrstrings);
843 Py_DECREF(attrstrings);
844 }
845 }
846 Py_ReprLeave((PyObject *)ag);
847 return repr;
848 }
849
850 static PyObject *
attrgetter_reduce(attrgetterobject * ag)851 attrgetter_reduce(attrgetterobject *ag)
852 {
853 PyObject *attrstrings = attrgetter_args(ag);
854 if (attrstrings == NULL)
855 return NULL;
856
857 return Py_BuildValue("ON", Py_TYPE(ag), attrstrings);
858 }
859
860 static PyMethodDef attrgetter_methods[] = {
861 {"__reduce__", (PyCFunction)attrgetter_reduce, METH_NOARGS,
862 reduce_doc},
863 {NULL}
864 };
865
866 PyDoc_STRVAR(attrgetter_doc,
867 "attrgetter(attr, ...) --> attrgetter object\n\
868 \n\
869 Return a callable object that fetches the given attribute(s) from its operand.\n\
870 After f = attrgetter('name'), the call f(r) returns r.name.\n\
871 After g = attrgetter('name', 'date'), the call g(r) returns (r.name, r.date).\n\
872 After h = attrgetter('name.first', 'name.last'), the call h(r) returns\n\
873 (r.name.first, r.name.last).");
874
875 static PyTypeObject attrgetter_type = {
876 PyVarObject_HEAD_INIT(NULL, 0)
877 "operator.attrgetter", /* tp_name */
878 sizeof(attrgetterobject), /* tp_basicsize */
879 0, /* tp_itemsize */
880 /* methods */
881 (destructor)attrgetter_dealloc, /* tp_dealloc */
882 0, /* tp_print */
883 0, /* tp_getattr */
884 0, /* tp_setattr */
885 0, /* tp_reserved */
886 (reprfunc)attrgetter_repr, /* tp_repr */
887 0, /* tp_as_number */
888 0, /* tp_as_sequence */
889 0, /* tp_as_mapping */
890 0, /* tp_hash */
891 (ternaryfunc)attrgetter_call, /* tp_call */
892 0, /* tp_str */
893 PyObject_GenericGetAttr, /* tp_getattro */
894 0, /* tp_setattro */
895 0, /* tp_as_buffer */
896 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
897 attrgetter_doc, /* tp_doc */
898 (traverseproc)attrgetter_traverse, /* tp_traverse */
899 0, /* tp_clear */
900 0, /* tp_richcompare */
901 0, /* tp_weaklistoffset */
902 0, /* tp_iter */
903 0, /* tp_iternext */
904 attrgetter_methods, /* tp_methods */
905 0, /* tp_members */
906 0, /* tp_getset */
907 0, /* tp_base */
908 0, /* tp_dict */
909 0, /* tp_descr_get */
910 0, /* tp_descr_set */
911 0, /* tp_dictoffset */
912 0, /* tp_init */
913 0, /* tp_alloc */
914 attrgetter_new, /* tp_new */
915 0, /* tp_free */
916 };
917
918
919 /* methodcaller object **********************************************************/
920
921 typedef struct {
922 PyObject_HEAD
923 PyObject *name;
924 PyObject *args;
925 PyObject *kwds;
926 } methodcallerobject;
927
928 static PyTypeObject methodcaller_type;
929
930 static PyObject *
methodcaller_new(PyTypeObject * type,PyObject * args,PyObject * kwds)931 methodcaller_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
932 {
933 methodcallerobject *mc;
934 PyObject *name;
935
936 if (PyTuple_GET_SIZE(args) < 1) {
937 PyErr_SetString(PyExc_TypeError, "methodcaller needs at least "
938 "one argument, the method name");
939 return NULL;
940 }
941
942 name = PyTuple_GET_ITEM(args, 0);
943 if (!PyUnicode_Check(name)) {
944 PyErr_SetString(PyExc_TypeError,
945 "method name must be a string");
946 return NULL;
947 }
948
949 /* create methodcallerobject structure */
950 mc = PyObject_GC_New(methodcallerobject, &methodcaller_type);
951 if (mc == NULL)
952 return NULL;
953
954 name = PyTuple_GET_ITEM(args, 0);
955 Py_INCREF(name);
956 PyUnicode_InternInPlace(&name);
957 mc->name = name;
958
959 Py_XINCREF(kwds);
960 mc->kwds = kwds;
961
962 mc->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
963 if (mc->args == NULL) {
964 Py_DECREF(mc);
965 return NULL;
966 }
967
968 PyObject_GC_Track(mc);
969 return (PyObject *)mc;
970 }
971
972 static void
methodcaller_dealloc(methodcallerobject * mc)973 methodcaller_dealloc(methodcallerobject *mc)
974 {
975 PyObject_GC_UnTrack(mc);
976 Py_XDECREF(mc->name);
977 Py_XDECREF(mc->args);
978 Py_XDECREF(mc->kwds);
979 PyObject_GC_Del(mc);
980 }
981
982 static int
methodcaller_traverse(methodcallerobject * mc,visitproc visit,void * arg)983 methodcaller_traverse(methodcallerobject *mc, visitproc visit, void *arg)
984 {
985 Py_VISIT(mc->args);
986 Py_VISIT(mc->kwds);
987 return 0;
988 }
989
990 static PyObject *
methodcaller_call(methodcallerobject * mc,PyObject * args,PyObject * kw)991 methodcaller_call(methodcallerobject *mc, PyObject *args, PyObject *kw)
992 {
993 PyObject *method, *obj, *result;
994
995 if (kw != NULL && !_PyArg_NoKeywords("methodcaller", kw))
996 return NULL;
997 if (!PyArg_UnpackTuple(args, "methodcaller", 1, 1, &obj))
998 return NULL;
999 method = PyObject_GetAttr(obj, mc->name);
1000 if (method == NULL)
1001 return NULL;
1002 result = PyObject_Call(method, mc->args, mc->kwds);
1003 Py_DECREF(method);
1004 return result;
1005 }
1006
1007 static PyObject *
methodcaller_repr(methodcallerobject * mc)1008 methodcaller_repr(methodcallerobject *mc)
1009 {
1010 PyObject *argreprs, *repr = NULL, *sep, *joinedargreprs;
1011 Py_ssize_t numtotalargs, numposargs, numkwdargs, i;
1012 int status = Py_ReprEnter((PyObject *)mc);
1013 if (status != 0) {
1014 if (status < 0)
1015 return NULL;
1016 return PyUnicode_FromFormat("%s(...)", Py_TYPE(mc)->tp_name);
1017 }
1018
1019 if (mc->kwds != NULL) {
1020 numkwdargs = PyDict_Size(mc->kwds);
1021 if (numkwdargs < 0) {
1022 Py_ReprLeave((PyObject *)mc);
1023 return NULL;
1024 }
1025 } else {
1026 numkwdargs = 0;
1027 }
1028
1029 numposargs = PyTuple_GET_SIZE(mc->args);
1030 numtotalargs = numposargs + numkwdargs;
1031
1032 if (numtotalargs == 0) {
1033 repr = PyUnicode_FromFormat("%s(%R)", Py_TYPE(mc)->tp_name, mc->name);
1034 Py_ReprLeave((PyObject *)mc);
1035 return repr;
1036 }
1037
1038 argreprs = PyTuple_New(numtotalargs);
1039 if (argreprs == NULL) {
1040 Py_ReprLeave((PyObject *)mc);
1041 return NULL;
1042 }
1043
1044 for (i = 0; i < numposargs; ++i) {
1045 PyObject *onerepr = PyObject_Repr(PyTuple_GET_ITEM(mc->args, i));
1046 if (onerepr == NULL)
1047 goto done;
1048 PyTuple_SET_ITEM(argreprs, i, onerepr);
1049 }
1050
1051 if (numkwdargs != 0) {
1052 PyObject *key, *value;
1053 Py_ssize_t pos = 0;
1054 while (PyDict_Next(mc->kwds, &pos, &key, &value)) {
1055 PyObject *onerepr = PyUnicode_FromFormat("%U=%R", key, value);
1056 if (onerepr == NULL)
1057 goto done;
1058 if (i >= numtotalargs) {
1059 i = -1;
1060 break;
1061 }
1062 PyTuple_SET_ITEM(argreprs, i, onerepr);
1063 ++i;
1064 }
1065 if (i != numtotalargs) {
1066 PyErr_SetString(PyExc_RuntimeError,
1067 "keywords dict changed size during iteration");
1068 goto done;
1069 }
1070 }
1071
1072 sep = PyUnicode_FromString(", ");
1073 if (sep == NULL)
1074 goto done;
1075
1076 joinedargreprs = PyUnicode_Join(sep, argreprs);
1077 Py_DECREF(sep);
1078 if (joinedargreprs == NULL)
1079 goto done;
1080
1081 repr = PyUnicode_FromFormat("%s(%R, %U)", Py_TYPE(mc)->tp_name,
1082 mc->name, joinedargreprs);
1083 Py_DECREF(joinedargreprs);
1084
1085 done:
1086 Py_DECREF(argreprs);
1087 Py_ReprLeave((PyObject *)mc);
1088 return repr;
1089 }
1090
1091 static PyObject *
methodcaller_reduce(methodcallerobject * mc)1092 methodcaller_reduce(methodcallerobject *mc)
1093 {
1094 PyObject *newargs;
1095 if (!mc->kwds || PyDict_Size(mc->kwds) == 0) {
1096 Py_ssize_t i;
1097 Py_ssize_t callargcount = PyTuple_GET_SIZE(mc->args);
1098 newargs = PyTuple_New(1 + callargcount);
1099 if (newargs == NULL)
1100 return NULL;
1101 Py_INCREF(mc->name);
1102 PyTuple_SET_ITEM(newargs, 0, mc->name);
1103 for (i = 0; i < callargcount; ++i) {
1104 PyObject *arg = PyTuple_GET_ITEM(mc->args, i);
1105 Py_INCREF(arg);
1106 PyTuple_SET_ITEM(newargs, i + 1, arg);
1107 }
1108 return Py_BuildValue("ON", Py_TYPE(mc), newargs);
1109 }
1110 else {
1111 PyObject *functools;
1112 PyObject *partial;
1113 PyObject *constructor;
1114 PyObject *newargs[2];
1115
1116 _Py_IDENTIFIER(partial);
1117 functools = PyImport_ImportModule("functools");
1118 if (!functools)
1119 return NULL;
1120 partial = _PyObject_GetAttrId(functools, &PyId_partial);
1121 Py_DECREF(functools);
1122 if (!partial)
1123 return NULL;
1124
1125 newargs[0] = (PyObject *)Py_TYPE(mc);
1126 newargs[1] = mc->name;
1127 constructor = _PyObject_FastCallDict(partial, newargs, 2, mc->kwds);
1128
1129 Py_DECREF(partial);
1130 return Py_BuildValue("NO", constructor, mc->args);
1131 }
1132 }
1133
1134 static PyMethodDef methodcaller_methods[] = {
1135 {"__reduce__", (PyCFunction)methodcaller_reduce, METH_NOARGS,
1136 reduce_doc},
1137 {NULL}
1138 };
1139 PyDoc_STRVAR(methodcaller_doc,
1140 "methodcaller(name, ...) --> methodcaller object\n\
1141 \n\
1142 Return a callable object that calls the given method on its operand.\n\
1143 After f = methodcaller('name'), the call f(r) returns r.name().\n\
1144 After g = methodcaller('name', 'date', foo=1), the call g(r) returns\n\
1145 r.name('date', foo=1).");
1146
1147 static PyTypeObject methodcaller_type = {
1148 PyVarObject_HEAD_INIT(NULL, 0)
1149 "operator.methodcaller", /* tp_name */
1150 sizeof(methodcallerobject), /* tp_basicsize */
1151 0, /* tp_itemsize */
1152 /* methods */
1153 (destructor)methodcaller_dealloc, /* tp_dealloc */
1154 0, /* tp_print */
1155 0, /* tp_getattr */
1156 0, /* tp_setattr */
1157 0, /* tp_reserved */
1158 (reprfunc)methodcaller_repr, /* tp_repr */
1159 0, /* tp_as_number */
1160 0, /* tp_as_sequence */
1161 0, /* tp_as_mapping */
1162 0, /* tp_hash */
1163 (ternaryfunc)methodcaller_call, /* tp_call */
1164 0, /* tp_str */
1165 PyObject_GenericGetAttr, /* tp_getattro */
1166 0, /* tp_setattro */
1167 0, /* tp_as_buffer */
1168 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
1169 methodcaller_doc, /* tp_doc */
1170 (traverseproc)methodcaller_traverse, /* tp_traverse */
1171 0, /* tp_clear */
1172 0, /* tp_richcompare */
1173 0, /* tp_weaklistoffset */
1174 0, /* tp_iter */
1175 0, /* tp_iternext */
1176 methodcaller_methods, /* tp_methods */
1177 0, /* tp_members */
1178 0, /* tp_getset */
1179 0, /* tp_base */
1180 0, /* tp_dict */
1181 0, /* tp_descr_get */
1182 0, /* tp_descr_set */
1183 0, /* tp_dictoffset */
1184 0, /* tp_init */
1185 0, /* tp_alloc */
1186 methodcaller_new, /* tp_new */
1187 0, /* tp_free */
1188 };
1189
1190
1191 /* Initialization function for the module (*must* be called PyInit__operator) */
1192
1193
1194 static struct PyModuleDef operatormodule = {
1195 PyModuleDef_HEAD_INIT,
1196 "_operator",
1197 operator_doc,
1198 -1,
1199 operator_methods,
1200 NULL,
1201 NULL,
1202 NULL,
1203 NULL
1204 };
1205
1206 PyMODINIT_FUNC
PyInit__operator(void)1207 PyInit__operator(void)
1208 {
1209 PyObject *m;
1210
1211 /* Create the module and add the functions */
1212 m = PyModule_Create(&operatormodule);
1213 if (m == NULL)
1214 return NULL;
1215
1216 if (PyType_Ready(&itemgetter_type) < 0)
1217 return NULL;
1218 Py_INCREF(&itemgetter_type);
1219 PyModule_AddObject(m, "itemgetter", (PyObject *)&itemgetter_type);
1220
1221 if (PyType_Ready(&attrgetter_type) < 0)
1222 return NULL;
1223 Py_INCREF(&attrgetter_type);
1224 PyModule_AddObject(m, "attrgetter", (PyObject *)&attrgetter_type);
1225
1226 if (PyType_Ready(&methodcaller_type) < 0)
1227 return NULL;
1228 Py_INCREF(&methodcaller_type);
1229 PyModule_AddObject(m, "methodcaller", (PyObject *)&methodcaller_type);
1230 return m;
1231 }
1232