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 PyInt_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 PyInt_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 
68 /* Deprecated operators that need warnings. */
69 static int
op_isCallable(PyObject * x)70 op_isCallable(PyObject *x)
71 {
72     if (PyErr_WarnPy3k("operator.isCallable() is not supported in 3.x. "
73                        "Use hasattr(obj, '__call__').", 1) < 0)
74         return -1;
75     return PyCallable_Check(x);
76 }
77 
78 static int
op_sequenceIncludes(PyObject * seq,PyObject * ob)79 op_sequenceIncludes(PyObject *seq, PyObject* ob)
80 {
81     if (PyErr_WarnPy3k("operator.sequenceIncludes() is not supported "
82                        "in 3.x. Use operator.contains().", 1) < 0)
83         return -1;
84     return PySequence_Contains(seq, ob);
85 }
86 
spami(isCallable,op_isCallable)87 spami(isCallable       , op_isCallable)
88 spami(isNumberType     , PyNumber_Check)
89 spami(truth            , PyObject_IsTrue)
90 spam2(op_add           , PyNumber_Add)
91 spam2(op_sub           , PyNumber_Subtract)
92 spam2(op_mul           , PyNumber_Multiply)
93 spam2(op_div           , PyNumber_Divide)
94 spam2(op_floordiv      , PyNumber_FloorDivide)
95 spam2(op_truediv       , PyNumber_TrueDivide)
96 spam2(op_mod           , PyNumber_Remainder)
97 spam1(op_neg           , PyNumber_Negative)
98 spam1(op_pos           , PyNumber_Positive)
99 spam1(op_abs           , PyNumber_Absolute)
100 spam1(op_inv           , PyNumber_Invert)
101 spam1(op_invert        , PyNumber_Invert)
102 spam2(op_lshift        , PyNumber_Lshift)
103 spam2(op_rshift        , PyNumber_Rshift)
104 spami(op_not_          , PyObject_Not)
105 spam2(op_and_          , PyNumber_And)
106 spam2(op_xor           , PyNumber_Xor)
107 spam2(op_or_           , PyNumber_Or)
108 spam2(op_iadd          , PyNumber_InPlaceAdd)
109 spam2(op_isub          , PyNumber_InPlaceSubtract)
110 spam2(op_imul          , PyNumber_InPlaceMultiply)
111 spam2(op_idiv          , PyNumber_InPlaceDivide)
112 spam2(op_ifloordiv     , PyNumber_InPlaceFloorDivide)
113 spam2(op_itruediv      , PyNumber_InPlaceTrueDivide)
114 spam2(op_imod          , PyNumber_InPlaceRemainder)
115 spam2(op_ilshift       , PyNumber_InPlaceLshift)
116 spam2(op_irshift       , PyNumber_InPlaceRshift)
117 spam2(op_iand          , PyNumber_InPlaceAnd)
118 spam2(op_ixor          , PyNumber_InPlaceXor)
119 spam2(op_ior           , PyNumber_InPlaceOr)
120 spami(isSequenceType   , PySequence_Check)
121 spam2(op_concat        , PySequence_Concat)
122 spamoi(op_repeat       , PySequence_Repeat)
123 spam2(op_iconcat       , PySequence_InPlaceConcat)
124 spamoi(op_irepeat      , PySequence_InPlaceRepeat)
125 spami2b(op_contains     , PySequence_Contains)
126 spami2b(sequenceIncludes, op_sequenceIncludes)
127 spamn2(indexOf         , PySequence_Index)
128 spamn2(countOf         , PySequence_Count)
129 spami(isMappingType    , PyMapping_Check)
130 spam2(op_getitem       , PyObject_GetItem)
131 spam2n(op_delitem       , PyObject_DelItem)
132 spam3n(op_setitem      , PyObject_SetItem)
133 spamrc(op_lt           , Py_LT)
134 spamrc(op_le           , Py_LE)
135 spamrc(op_eq           , Py_EQ)
136 spamrc(op_ne           , Py_NE)
137 spamrc(op_gt           , Py_GT)
138 spamrc(op_ge           , Py_GE)
139 
140 static PyObject*
141 op_pow(PyObject *s, PyObject *a)
142 {
143     PyObject *a1, *a2;
144     if (PyArg_UnpackTuple(a,"pow", 2, 2, &a1, &a2))
145         return PyNumber_Power(a1, a2, Py_None);
146     return NULL;
147 }
148 
149 static PyObject*
op_ipow(PyObject * s,PyObject * a)150 op_ipow(PyObject *s, PyObject *a)
151 {
152     PyObject *a1, *a2;
153     if (PyArg_UnpackTuple(a,"ipow", 2, 2, &a1, &a2))
154         return PyNumber_InPlacePower(a1, a2, Py_None);
155     return NULL;
156 }
157 
158 static PyObject *
op_index(PyObject * s,PyObject * a)159 op_index(PyObject *s, PyObject *a)
160 {
161     return PyNumber_Index(a);
162 }
163 
164 static PyObject*
is_(PyObject * s,PyObject * a)165 is_(PyObject *s, PyObject *a)
166 {
167     PyObject *a1, *a2, *result = NULL;
168     if (PyArg_UnpackTuple(a,"is_", 2, 2, &a1, &a2)) {
169         result = (a1 == a2) ? Py_True : Py_False;
170         Py_INCREF(result);
171     }
172     return result;
173 }
174 
175 static PyObject*
is_not(PyObject * s,PyObject * a)176 is_not(PyObject *s, PyObject *a)
177 {
178     PyObject *a1, *a2, *result = NULL;
179     if (PyArg_UnpackTuple(a,"is_not", 2, 2, &a1, &a2)) {
180         result = (a1 != a2) ? Py_True : Py_False;
181         Py_INCREF(result);
182     }
183     return result;
184 }
185 
186 static PyObject*
op_getslice(PyObject * s,PyObject * a)187 op_getslice(PyObject *s, PyObject *a)
188 {
189     PyObject *a1;
190     Py_ssize_t a2, a3;
191 
192     if (!PyArg_ParseTuple(a, "Onn:getslice", &a1, &a2, &a3))
193         return NULL;
194     return PySequence_GetSlice(a1, a2, a3);
195 }
196 
197 static PyObject*
op_setslice(PyObject * s,PyObject * a)198 op_setslice(PyObject *s, PyObject *a)
199 {
200     PyObject *a1, *a4;
201     Py_ssize_t a2, a3;
202 
203     if (!PyArg_ParseTuple(a, "OnnO:setslice", &a1, &a2, &a3, &a4))
204         return NULL;
205 
206     if (-1 == PySequence_SetSlice(a1, a2, a3, a4))
207         return NULL;
208 
209     Py_RETURN_NONE;
210 }
211 
212 static PyObject*
op_delslice(PyObject * s,PyObject * a)213 op_delslice(PyObject *s, PyObject *a)
214 {
215     PyObject *a1;
216     Py_ssize_t a2, a3;
217 
218     if (!PyArg_ParseTuple(a, "Onn:delslice", &a1, &a2, &a3))
219         return NULL;
220 
221     if (-1 == PySequence_DelSlice(a1, a2, a3))
222         return NULL;
223 
224     Py_RETURN_NONE;
225 }
226 
227 #undef spam1
228 #undef spam2
229 #undef spam1o
230 #undef spam1o
231 #define spam1(OP,DOC) {#OP, OP, METH_VARARGS, PyDoc_STR(DOC)},
232 #define spam2(OP,ALTOP,DOC) {#OP, op_##OP, METH_VARARGS, PyDoc_STR(DOC)}, \
233                            {#ALTOP, op_##OP, METH_VARARGS, PyDoc_STR(DOC)},
234 #define spam1o(OP,DOC) {#OP, OP, METH_O, PyDoc_STR(DOC)},
235 #define spam2o(OP,ALTOP,DOC) {#OP, op_##OP, METH_O, PyDoc_STR(DOC)}, \
236                            {#ALTOP, op_##OP, METH_O, PyDoc_STR(DOC)},
237 
238 
239 
240 /* compare_digest **********************************************************/
241 
242 /*
243  * timing safe compare
244  *
245  * Returns 1 of the strings are equal.
246  * In case of len(a) != len(b) the function tries to keep the timing
247  * dependent on the length of b. CPU cache locally may still alter timing
248  * a bit.
249  */
250 static int
_tscmp(const unsigned char * a,const unsigned char * b,Py_ssize_t len_a,Py_ssize_t len_b)251 _tscmp(const unsigned char *a, const unsigned char *b,
252         Py_ssize_t len_a, Py_ssize_t len_b)
253 {
254     /* The volatile type declarations make sure that the compiler has no
255      * chance to optimize and fold the code in any way that may change
256      * the timing.
257      */
258     volatile Py_ssize_t length;
259     volatile const unsigned char *left;
260     volatile const unsigned char *right;
261     Py_ssize_t i;
262     unsigned char result;
263 
264     /* loop count depends on length of b */
265     length = len_b;
266     left = NULL;
267     right = b;
268 
269     /* don't use else here to keep the amount of CPU instructions constant,
270      * volatile forces re-evaluation
271      *  */
272     if (len_a == length) {
273         left = *((volatile const unsigned char**)&a);
274         result = 0;
275     }
276     if (len_a != length) {
277         left = b;
278         result = 1;
279     }
280 
281     for (i=0; i < length; i++) {
282         result |= *left++ ^ *right++;
283     }
284 
285     return (result == 0);
286 }
287 
288 PyDoc_STRVAR(compare_digest__doc__,
289 "compare_digest(a, b) -> bool\n"
290 "\n"
291 "Return 'a == b'.  This function uses an approach designed to prevent\n"
292 "timing analysis, making it appropriate for cryptography.\n"
293 "a and b must both be of the same type: either str (ASCII only),\n"
294 "or any type that supports the buffer protocol (e.g. bytes).\n"
295 "\n"
296 "Note: If a and b are of different lengths, or if an error occurs,\n"
297 "a timing attack could theoretically reveal information about the\n"
298 "types and lengths of a and b--but not their values.\n");
299 
300 static PyObject*
compare_digest(PyObject * self,PyObject * args)301 compare_digest(PyObject *self, PyObject *args)
302 {
303     PyObject *a, *b;
304     int rc;
305 
306     if (!PyArg_ParseTuple(args, "OO:compare_digest", &a, &b)) {
307         return NULL;
308     }
309 
310     /* Unicode string */
311     if (PyUnicode_Check(a) && PyUnicode_Check(b)) {
312         rc = _tscmp((const unsigned char *)PyUnicode_AS_DATA(a),
313                     (const unsigned char *)PyUnicode_AS_DATA(b),
314                     PyUnicode_GET_DATA_SIZE(a),
315                     PyUnicode_GET_DATA_SIZE(b));
316     }
317     /* fallback to buffer interface for bytes, bytesarray and other */
318     else {
319         Py_buffer view_a;
320         Py_buffer view_b;
321 
322         if (PyObject_CheckBuffer(a) == 0 && PyObject_CheckBuffer(b) == 0) {
323             PyErr_Format(PyExc_TypeError,
324                          "unsupported operand types(s) or combination of types: "
325                          "'%.100s' and '%.100s'",
326                          Py_TYPE(a)->tp_name, Py_TYPE(b)->tp_name);
327             return NULL;
328         }
329 
330         if (PyObject_GetBuffer(a, &view_a, PyBUF_SIMPLE) == -1) {
331             return NULL;
332         }
333         if (view_a.ndim > 1) {
334             PyErr_SetString(PyExc_BufferError,
335                             "Buffer must be single dimension");
336             PyBuffer_Release(&view_a);
337             return NULL;
338         }
339 
340         if (PyObject_GetBuffer(b, &view_b, PyBUF_SIMPLE) == -1) {
341             PyBuffer_Release(&view_a);
342             return NULL;
343         }
344         if (view_b.ndim > 1) {
345             PyErr_SetString(PyExc_BufferError,
346                             "Buffer must be single dimension");
347             PyBuffer_Release(&view_a);
348             PyBuffer_Release(&view_b);
349             return NULL;
350         }
351 
352         rc = _tscmp((const unsigned char*)view_a.buf,
353                     (const unsigned char*)view_b.buf,
354                     view_a.len,
355                     view_b.len);
356 
357         PyBuffer_Release(&view_a);
358         PyBuffer_Release(&view_b);
359     }
360 
361     return PyBool_FromLong(rc);
362 }
363 
364 static struct PyMethodDef operator_methods[] = {
365 
366 spam1o(isCallable,
367  "isCallable(a) -- Same as callable(a).")
368 spam1o(isNumberType,
369  "isNumberType(a) -- Return True if a has a numeric type, False otherwise.")
370 spam1o(isSequenceType,
371  "isSequenceType(a) -- Return True if a has a sequence type, False otherwise.")
372 spam1o(truth,
373  "truth(a) -- Return True if a is true, False otherwise.")
374 spam2(contains,__contains__,
375  "contains(a, b) -- Same as b in a (note reversed operands).")
376 spam1(sequenceIncludes,
377  "sequenceIncludes(a, b) -- Same as b in a (note reversed operands; deprecated).")
378 spam1(indexOf,
379  "indexOf(a, b) -- Return the first index of b in a.")
380 spam1(countOf,
381  "countOf(a, b) -- Return the number of times b occurs in a.")
382 spam1o(isMappingType,
383  "isMappingType(a) -- Return True if a has a mapping type, False otherwise.")
384 
385 spam1(is_, "is_(a, b) -- Same as a is b.")
386 spam1(is_not, "is_not(a, b) -- Same as a is not b.")
387 spam2o(index, __index__, "index(a) -- Same as a.__index__()")
388 spam2(add,__add__, "add(a, b) -- Same as a + b.")
389 spam2(sub,__sub__, "sub(a, b) -- Same as a - b.")
390 spam2(mul,__mul__, "mul(a, b) -- Same as a * b.")
391 spam2(div,__div__, "div(a, b) -- Same as a / b when __future__.division is not in effect.")
392 spam2(floordiv,__floordiv__, "floordiv(a, b) -- Same as a // b.")
393 spam2(truediv,__truediv__, "truediv(a, b) -- Same as a / b when __future__.division is in effect.")
394 spam2(mod,__mod__, "mod(a, b) -- Same as a % b.")
395 spam2o(neg,__neg__, "neg(a) -- Same as -a.")
396 spam2o(pos,__pos__, "pos(a) -- Same as +a.")
397 spam2o(abs,__abs__, "abs(a) -- Same as abs(a).")
398 spam2o(inv,__inv__, "inv(a) -- Same as ~a.")
399 spam2o(invert,__invert__, "invert(a) -- Same as ~a.")
400 spam2(lshift,__lshift__, "lshift(a, b) -- Same as a << b.")
401 spam2(rshift,__rshift__, "rshift(a, b) -- Same as a >> b.")
402 spam2o(not_,__not__, "not_(a) -- Same as not a.")
403 spam2(and_,__and__, "and_(a, b) -- Same as a & b.")
404 spam2(xor,__xor__, "xor(a, b) -- Same as a ^ b.")
405 spam2(or_,__or__, "or_(a, b) -- Same as a | b.")
406 spam2(iadd,__iadd__, "a = iadd(a, b) -- Same as a += b.")
407 spam2(isub,__isub__, "a = isub(a, b) -- Same as a -= b.")
408 spam2(imul,__imul__, "a = imul(a, b) -- Same as a *= b.")
409 spam2(idiv,__idiv__, "a = idiv(a, b) -- Same as a /= b when __future__.division is not in effect.")
410 spam2(ifloordiv,__ifloordiv__, "a = ifloordiv(a, b) -- Same as a //= b.")
411 spam2(itruediv,__itruediv__, "a = itruediv(a, b) -- Same as a /= b when __future__.division is in effect.")
412 spam2(imod,__imod__, "a = imod(a, b) -- Same as a %= b.")
413 spam2(ilshift,__ilshift__, "a = ilshift(a, b) -- Same as a <<= b.")
414 spam2(irshift,__irshift__, "a = irshift(a, b) -- Same as a >>= b.")
415 spam2(iand,__iand__, "a = iand(a, b) -- Same as a &= b.")
416 spam2(ixor,__ixor__, "a = ixor(a, b) -- Same as a ^= b.")
417 spam2(ior,__ior__, "a = ior(a, b) -- Same as a |= b.")
418 spam2(concat,__concat__,
419  "concat(a, b) -- Same as a + b, for a and b sequences.")
420 spam2(repeat,__repeat__,
421  "repeat(a, b) -- Return a * b, where a is a sequence, and b is an integer.")
422 spam2(iconcat,__iconcat__,
423  "a = iconcat(a, b) -- Same as a += b, for a and b sequences.")
424 spam2(irepeat,__irepeat__,
425  "a = irepeat(a, b) -- Same as a *= b, where a is a sequence, and b is an integer.")
426 spam2(getitem,__getitem__,
427  "getitem(a, b) -- Same as a[b].")
428 spam2(setitem,__setitem__,
429  "setitem(a, b, c) -- Same as a[b] = c.")
430 spam2(delitem,__delitem__,
431  "delitem(a, b) -- Same as del a[b].")
432 spam2(pow,__pow__, "pow(a, b) -- Same as a ** b.")
433 spam2(ipow,__ipow__, "a = ipow(a, b) -- Same as a **= b.")
434 spam2(getslice,__getslice__,
435  "getslice(a, b, c) -- Same as a[b:c].")
436 spam2(setslice,__setslice__,
437 "setslice(a, b, c, d) -- Same as a[b:c] = d.")
438 spam2(delslice,__delslice__,
439 "delslice(a, b, c) -- Same as del a[b:c].")
440 spam2(lt,__lt__, "lt(a, b) -- Same as a<b.")
441 spam2(le,__le__, "le(a, b) -- Same as a<=b.")
442 spam2(eq,__eq__, "eq(a, b) -- Same as a==b.")
443 spam2(ne,__ne__, "ne(a, b) -- Same as a!=b.")
444 spam2(gt,__gt__, "gt(a, b) -- Same as a>b.")
445 spam2(ge,__ge__, "ge(a, b) -- Same as a>=b.")
446 
447     {"_compare_digest", (PyCFunction)compare_digest, METH_VARARGS,
448      compare_digest__doc__},
449     {NULL,              NULL}           /* sentinel */
450 
451 };
452 
453 /* itemgetter object **********************************************************/
454 
455 typedef struct {
456     PyObject_HEAD
457     Py_ssize_t nitems;
458     PyObject *item;
459 } itemgetterobject;
460 
461 static PyTypeObject itemgetter_type;
462 
463 static PyObject *
itemgetter_new(PyTypeObject * type,PyObject * args,PyObject * kwds)464 itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
465 {
466     itemgetterobject *ig;
467     PyObject *item;
468     Py_ssize_t nitems;
469 
470     if (!_PyArg_NoKeywords("itemgetter()", kwds))
471         return NULL;
472 
473     nitems = PyTuple_GET_SIZE(args);
474     if (nitems <= 1) {
475         if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &item))
476             return NULL;
477     } else
478         item = args;
479 
480     /* create itemgetterobject structure */
481     ig = PyObject_GC_New(itemgetterobject, &itemgetter_type);
482     if (ig == NULL)
483         return NULL;
484 
485     Py_INCREF(item);
486     ig->item = item;
487     ig->nitems = nitems;
488 
489     PyObject_GC_Track(ig);
490     return (PyObject *)ig;
491 }
492 
493 static void
itemgetter_dealloc(itemgetterobject * ig)494 itemgetter_dealloc(itemgetterobject *ig)
495 {
496     PyObject_GC_UnTrack(ig);
497     Py_XDECREF(ig->item);
498     PyObject_GC_Del(ig);
499 }
500 
501 static int
itemgetter_traverse(itemgetterobject * ig,visitproc visit,void * arg)502 itemgetter_traverse(itemgetterobject *ig, visitproc visit, void *arg)
503 {
504     Py_VISIT(ig->item);
505     return 0;
506 }
507 
508 static PyObject *
itemgetter_call(itemgetterobject * ig,PyObject * args,PyObject * kw)509 itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw)
510 {
511     PyObject *obj, *result;
512     Py_ssize_t i, nitems=ig->nitems;
513 
514     if (kw != NULL && !_PyArg_NoKeywords("itemgetter", kw))
515         return NULL;
516     if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj))
517         return NULL;
518     if (nitems == 1)
519         return PyObject_GetItem(obj, ig->item);
520 
521     assert(PyTuple_Check(ig->item));
522     assert(PyTuple_GET_SIZE(ig->item) == nitems);
523 
524     result = PyTuple_New(nitems);
525     if (result == NULL)
526         return NULL;
527 
528     for (i=0 ; i < nitems ; i++) {
529         PyObject *item, *val;
530         item = PyTuple_GET_ITEM(ig->item, i);
531         val = PyObject_GetItem(obj, item);
532         if (val == NULL) {
533             Py_DECREF(result);
534             return NULL;
535         }
536         PyTuple_SET_ITEM(result, i, val);
537     }
538     return result;
539 }
540 
541 PyDoc_STRVAR(itemgetter_doc,
542 "itemgetter(item, ...) --> itemgetter object\n\
543 \n\
544 Return a callable object that fetches the given item(s) from its operand.\n\
545 After f = itemgetter(2), the call f(r) returns r[2].\n\
546 After g = itemgetter(2, 5, 3), the call g(r) returns (r[2], r[5], r[3])");
547 
548 static PyTypeObject itemgetter_type = {
549     PyVarObject_HEAD_INIT(NULL, 0)
550     "operator.itemgetter",              /* tp_name */
551     sizeof(itemgetterobject),           /* tp_basicsize */
552     0,                                  /* tp_itemsize */
553     /* methods */
554     (destructor)itemgetter_dealloc,     /* tp_dealloc */
555     0,                                  /* tp_print */
556     0,                                  /* tp_getattr */
557     0,                                  /* tp_setattr */
558     0,                                  /* tp_compare */
559     0,                                  /* tp_repr */
560     0,                                  /* tp_as_number */
561     0,                                  /* tp_as_sequence */
562     0,                                  /* tp_as_mapping */
563     0,                                  /* tp_hash */
564     (ternaryfunc)itemgetter_call,       /* tp_call */
565     0,                                  /* tp_str */
566     PyObject_GenericGetAttr,            /* tp_getattro */
567     0,                                  /* tp_setattro */
568     0,                                  /* tp_as_buffer */
569     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,            /* tp_flags */
570     itemgetter_doc,                     /* tp_doc */
571     (traverseproc)itemgetter_traverse,          /* tp_traverse */
572     0,                                  /* tp_clear */
573     0,                                  /* tp_richcompare */
574     0,                                  /* tp_weaklistoffset */
575     0,                                  /* tp_iter */
576     0,                                  /* tp_iternext */
577     0,                                  /* tp_methods */
578     0,                                  /* tp_members */
579     0,                                  /* tp_getset */
580     0,                                  /* tp_base */
581     0,                                  /* tp_dict */
582     0,                                  /* tp_descr_get */
583     0,                                  /* tp_descr_set */
584     0,                                  /* tp_dictoffset */
585     0,                                  /* tp_init */
586     0,                                  /* tp_alloc */
587     itemgetter_new,                     /* tp_new */
588     0,                                  /* tp_free */
589 };
590 
591 
592 /* attrgetter object **********************************************************/
593 
594 typedef struct {
595     PyObject_HEAD
596     Py_ssize_t nattrs;
597     PyObject *attr;
598 } attrgetterobject;
599 
600 static PyTypeObject attrgetter_type;
601 
602 static PyObject *
attrgetter_new(PyTypeObject * type,PyObject * args,PyObject * kwds)603 attrgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
604 {
605     attrgetterobject *ag;
606     PyObject *attr;
607     Py_ssize_t nattrs;
608 
609     if (!_PyArg_NoKeywords("attrgetter()", kwds))
610         return NULL;
611 
612     nattrs = PyTuple_GET_SIZE(args);
613     if (nattrs <= 1) {
614         if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &attr))
615             return NULL;
616     } else
617         attr = args;
618 
619     /* create attrgetterobject structure */
620     ag = PyObject_GC_New(attrgetterobject, &attrgetter_type);
621     if (ag == NULL)
622         return NULL;
623 
624     Py_INCREF(attr);
625     ag->attr = attr;
626     ag->nattrs = nattrs;
627 
628     PyObject_GC_Track(ag);
629     return (PyObject *)ag;
630 }
631 
632 static void
attrgetter_dealloc(attrgetterobject * ag)633 attrgetter_dealloc(attrgetterobject *ag)
634 {
635     PyObject_GC_UnTrack(ag);
636     Py_XDECREF(ag->attr);
637     PyObject_GC_Del(ag);
638 }
639 
640 static int
attrgetter_traverse(attrgetterobject * ag,visitproc visit,void * arg)641 attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg)
642 {
643     Py_VISIT(ag->attr);
644     return 0;
645 }
646 
647 static PyObject *
dotted_getattr(PyObject * obj,PyObject * attr)648 dotted_getattr(PyObject *obj, PyObject *attr)
649 {
650     char *s, *p;
651 
652 #ifdef Py_USING_UNICODE
653     if (PyUnicode_Check(attr)) {
654         attr = _PyUnicode_AsDefaultEncodedString(attr, NULL);
655         if (attr == NULL)
656             return NULL;
657     }
658 #endif
659 
660     if (!PyString_Check(attr)) {
661         PyErr_SetString(PyExc_TypeError,
662                         "attribute name must be a string");
663         return NULL;
664     }
665 
666     s = PyString_AS_STRING(attr);
667     Py_INCREF(obj);
668     for (;;) {
669         PyObject *newobj, *str;
670         p = strchr(s, '.');
671         str = p ? PyString_FromStringAndSize(s, (p-s)) :
672               PyString_FromString(s);
673         if (str == NULL) {
674             Py_DECREF(obj);
675             return NULL;
676         }
677         newobj = PyObject_GetAttr(obj, str);
678         Py_DECREF(str);
679         Py_DECREF(obj);
680         if (newobj == NULL)
681             return NULL;
682         obj = newobj;
683         if (p == NULL) break;
684         s = p+1;
685     }
686 
687     return obj;
688 }
689 
690 static PyObject *
attrgetter_call(attrgetterobject * ag,PyObject * args,PyObject * kw)691 attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
692 {
693     PyObject *obj, *result;
694     Py_ssize_t i, nattrs=ag->nattrs;
695 
696     if (kw != NULL && !_PyArg_NoKeywords("attrgetter", kw))
697         return NULL;
698     if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj))
699         return NULL;
700     if (ag->nattrs == 1)
701         return dotted_getattr(obj, ag->attr);
702 
703     assert(PyTuple_Check(ag->attr));
704     assert(PyTuple_GET_SIZE(ag->attr) == nattrs);
705 
706     result = PyTuple_New(nattrs);
707     if (result == NULL)
708         return NULL;
709 
710     for (i=0 ; i < nattrs ; i++) {
711         PyObject *attr, *val;
712         attr = PyTuple_GET_ITEM(ag->attr, i);
713         val = dotted_getattr(obj, attr);
714         if (val == NULL) {
715             Py_DECREF(result);
716             return NULL;
717         }
718         PyTuple_SET_ITEM(result, i, val);
719     }
720     return result;
721 }
722 
723 PyDoc_STRVAR(attrgetter_doc,
724 "attrgetter(attr, ...) --> attrgetter object\n\
725 \n\
726 Return a callable object that fetches the given attribute(s) from its operand.\n\
727 After f = attrgetter('name'), the call f(r) returns r.name.\n\
728 After g = attrgetter('name', 'date'), the call g(r) returns (r.name, r.date).\n\
729 After h = attrgetter('name.first', 'name.last'), the call h(r) returns\n\
730 (r.name.first, r.name.last).");
731 
732 static PyTypeObject attrgetter_type = {
733     PyVarObject_HEAD_INIT(NULL, 0)
734     "operator.attrgetter",              /* tp_name */
735     sizeof(attrgetterobject),           /* tp_basicsize */
736     0,                                  /* tp_itemsize */
737     /* methods */
738     (destructor)attrgetter_dealloc,     /* tp_dealloc */
739     0,                                  /* tp_print */
740     0,                                  /* tp_getattr */
741     0,                                  /* tp_setattr */
742     0,                                  /* tp_compare */
743     0,                                  /* tp_repr */
744     0,                                  /* tp_as_number */
745     0,                                  /* tp_as_sequence */
746     0,                                  /* tp_as_mapping */
747     0,                                  /* tp_hash */
748     (ternaryfunc)attrgetter_call,       /* tp_call */
749     0,                                  /* tp_str */
750     PyObject_GenericGetAttr,            /* tp_getattro */
751     0,                                  /* tp_setattro */
752     0,                                  /* tp_as_buffer */
753     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,            /* tp_flags */
754     attrgetter_doc,                     /* tp_doc */
755     (traverseproc)attrgetter_traverse,          /* tp_traverse */
756     0,                                  /* tp_clear */
757     0,                                  /* tp_richcompare */
758     0,                                  /* tp_weaklistoffset */
759     0,                                  /* tp_iter */
760     0,                                  /* tp_iternext */
761     0,                                  /* tp_methods */
762     0,                                  /* tp_members */
763     0,                                  /* tp_getset */
764     0,                                  /* tp_base */
765     0,                                  /* tp_dict */
766     0,                                  /* tp_descr_get */
767     0,                                  /* tp_descr_set */
768     0,                                  /* tp_dictoffset */
769     0,                                  /* tp_init */
770     0,                                  /* tp_alloc */
771     attrgetter_new,                     /* tp_new */
772     0,                                  /* tp_free */
773 };
774 
775 
776 /* methodcaller object **********************************************************/
777 
778 typedef struct {
779     PyObject_HEAD
780     PyObject *name;
781     PyObject *args;
782     PyObject *kwds;
783 } methodcallerobject;
784 
785 static PyTypeObject methodcaller_type;
786 
787 static PyObject *
methodcaller_new(PyTypeObject * type,PyObject * args,PyObject * kwds)788 methodcaller_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
789 {
790     methodcallerobject *mc;
791     PyObject *name;
792 
793     if (PyTuple_GET_SIZE(args) < 1) {
794         PyErr_SetString(PyExc_TypeError, "methodcaller needs at least "
795                         "one argument, the method name");
796         return NULL;
797     }
798 
799     /* create methodcallerobject structure */
800     mc = PyObject_GC_New(methodcallerobject, &methodcaller_type);
801     if (mc == NULL)
802         return NULL;
803 
804     name = PyTuple_GET_ITEM(args, 0);
805     Py_INCREF(name);
806     mc->name = name;
807 
808     Py_XINCREF(kwds);
809     mc->kwds = kwds;
810 
811     mc->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
812     if (mc->args == NULL) {
813         Py_DECREF(mc);
814         return NULL;
815     }
816 
817     PyObject_GC_Track(mc);
818     return (PyObject *)mc;
819 }
820 
821 static void
methodcaller_dealloc(methodcallerobject * mc)822 methodcaller_dealloc(methodcallerobject *mc)
823 {
824     PyObject_GC_UnTrack(mc);
825     Py_XDECREF(mc->name);
826     Py_XDECREF(mc->args);
827     Py_XDECREF(mc->kwds);
828     PyObject_GC_Del(mc);
829 }
830 
831 static int
methodcaller_traverse(methodcallerobject * mc,visitproc visit,void * arg)832 methodcaller_traverse(methodcallerobject *mc, visitproc visit, void *arg)
833 {
834     Py_VISIT(mc->args);
835     Py_VISIT(mc->kwds);
836     return 0;
837 }
838 
839 static PyObject *
methodcaller_call(methodcallerobject * mc,PyObject * args,PyObject * kw)840 methodcaller_call(methodcallerobject *mc, PyObject *args, PyObject *kw)
841 {
842     PyObject *method, *obj, *result;
843 
844     if (kw != NULL && !_PyArg_NoKeywords("methodcaller", kw))
845         return NULL;
846     if (!PyArg_UnpackTuple(args, "methodcaller", 1, 1, &obj))
847         return NULL;
848     method = PyObject_GetAttr(obj, mc->name);
849     if (method == NULL)
850         return NULL;
851     result = PyObject_Call(method, mc->args, mc->kwds);
852     Py_DECREF(method);
853     return result;
854 }
855 
856 PyDoc_STRVAR(methodcaller_doc,
857 "methodcaller(name, ...) --> methodcaller object\n\
858 \n\
859 Return a callable object that calls the given method on its operand.\n\
860 After f = methodcaller('name'), the call f(r) returns r.name().\n\
861 After g = methodcaller('name', 'date', foo=1), the call g(r) returns\n\
862 r.name('date', foo=1).");
863 
864 static PyTypeObject methodcaller_type = {
865     PyVarObject_HEAD_INIT(NULL, 0)
866     "operator.methodcaller",            /* tp_name */
867     sizeof(methodcallerobject),         /* tp_basicsize */
868     0,                                  /* tp_itemsize */
869     /* methods */
870     (destructor)methodcaller_dealloc, /* tp_dealloc */
871     0,                                  /* tp_print */
872     0,                                  /* tp_getattr */
873     0,                                  /* tp_setattr */
874     0,                                  /* tp_compare */
875     0,                                  /* tp_repr */
876     0,                                  /* tp_as_number */
877     0,                                  /* tp_as_sequence */
878     0,                                  /* tp_as_mapping */
879     0,                                  /* tp_hash */
880     (ternaryfunc)methodcaller_call,     /* tp_call */
881     0,                                  /* tp_str */
882     PyObject_GenericGetAttr,            /* tp_getattro */
883     0,                                  /* tp_setattro */
884     0,                                  /* tp_as_buffer */
885     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
886     methodcaller_doc,                           /* tp_doc */
887     (traverseproc)methodcaller_traverse,        /* tp_traverse */
888     0,                                  /* tp_clear */
889     0,                                  /* tp_richcompare */
890     0,                                  /* tp_weaklistoffset */
891     0,                                  /* tp_iter */
892     0,                                  /* tp_iternext */
893     0,                                  /* tp_methods */
894     0,                                  /* tp_members */
895     0,                                  /* tp_getset */
896     0,                                  /* tp_base */
897     0,                                  /* tp_dict */
898     0,                                  /* tp_descr_get */
899     0,                                  /* tp_descr_set */
900     0,                                  /* tp_dictoffset */
901     0,                                  /* tp_init */
902     0,                                  /* tp_alloc */
903     methodcaller_new,                   /* tp_new */
904     0,                                  /* tp_free */
905 };
906 
907 
908 /* Initialization function for the module (*must* be called initoperator) */
909 
910 PyMODINIT_FUNC
initoperator(void)911 initoperator(void)
912 {
913     PyObject *m;
914 
915     /* Create the module and add the functions */
916     m = Py_InitModule4("operator", operator_methods, operator_doc,
917                    (PyObject*)NULL, PYTHON_API_VERSION);
918     if (m == NULL)
919         return;
920 
921     if (PyType_Ready(&itemgetter_type) < 0)
922         return;
923     Py_INCREF(&itemgetter_type);
924     PyModule_AddObject(m, "itemgetter", (PyObject *)&itemgetter_type);
925 
926     if (PyType_Ready(&attrgetter_type) < 0)
927         return;
928     Py_INCREF(&attrgetter_type);
929     PyModule_AddObject(m, "attrgetter", (PyObject *)&attrgetter_type);
930 
931     if (PyType_Ready(&methodcaller_type) < 0)
932         return;
933     Py_INCREF(&methodcaller_type);
934     PyModule_AddObject(m, "methodcaller", (PyObject *)&methodcaller_type);
935 }
936