blob: 8ca4d2d06cdb87d065700969d4ea27b009aa3a4f [file] [log] [blame]
Martin v. Löwis14f8b4c2002-06-13 20:33:02 +00001
2#include "Python.h"
3
4PyDoc_STRVAR(operator_doc,
5"Operator interface.\n\
Guido van Rossum037b9401996-07-30 16:55:54 +00006\n\
7This module exports a set of functions implemented in C corresponding\n\
8to the intrinsic operators of Python. For example, operator.add(x, y)\n\
9is equivalent to the expression x+y. The function names are those\n\
Benjamin Petersona0dfa822009-11-13 02:25:08 +000010used for special methods; variants without leading and trailing\n\
Martin v. Löwis14f8b4c2002-06-13 20:33:02 +000011'__' are also provided for convenience.");
Guido van Rossum037b9401996-07-30 16:55:54 +000012
Raymond Hettinger36cd2bf2003-01-03 08:24:58 +000013#define spam1(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a1) { \
Guido van Rossum037b9401996-07-30 16:55:54 +000014 return AOP(a1); }
15
Fred Drake5639ba42000-07-08 04:12:08 +000016#define spam2(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
Guido van Rossum037b9401996-07-30 16:55:54 +000017 PyObject *a1, *a2; \
Raymond Hettingerea3fdf42002-12-29 16:33:45 +000018 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
Guido van Rossum037b9401996-07-30 16:55:54 +000019 return AOP(a1,a2); }
20
Fred Drake5639ba42000-07-08 04:12:08 +000021#define spamoi(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
Guido van Rossum037b9401996-07-30 16:55:54 +000022 PyObject *a1; int a2; \
Fred Drakeea4d3f02000-09-17 16:09:27 +000023 if(! PyArg_ParseTuple(a,"Oi:" #OP,&a1,&a2)) return NULL; \
Guido van Rossum037b9401996-07-30 16:55:54 +000024 return AOP(a1,a2); }
25
Fred Drake5639ba42000-07-08 04:12:08 +000026#define spam2n(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
Guido van Rossum037b9401996-07-30 16:55:54 +000027 PyObject *a1, *a2; \
Raymond Hettingerea3fdf42002-12-29 16:33:45 +000028 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
Guido van Rossum037b9401996-07-30 16:55:54 +000029 if(-1 == AOP(a1,a2)) return NULL; \
30 Py_INCREF(Py_None); \
31 return Py_None; }
32
Fred Drake5639ba42000-07-08 04:12:08 +000033#define spam3n(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
Guido van Rossum037b9401996-07-30 16:55:54 +000034 PyObject *a1, *a2, *a3; \
Raymond Hettingerea3fdf42002-12-29 16:33:45 +000035 if(! PyArg_UnpackTuple(a,#OP,3,3,&a1,&a2,&a3)) return NULL; \
Guido van Rossum037b9401996-07-30 16:55:54 +000036 if(-1 == AOP(a1,a2,a3)) return NULL; \
37 Py_INCREF(Py_None); \
38 return Py_None; }
39
Raymond Hettinger36cd2bf2003-01-03 08:24:58 +000040#define spami(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a1) { \
41 long r; \
Guido van Rossum037b9401996-07-30 16:55:54 +000042 if(-1 == (r=AOP(a1))) return NULL; \
Guido van Rossum77f6a652002-04-03 22:41:51 +000043 return PyBool_FromLong(r); }
Guido van Rossum037b9401996-07-30 16:55:54 +000044
Fred Drake5639ba42000-07-08 04:12:08 +000045#define spami2(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
Guido van Rossum037b9401996-07-30 16:55:54 +000046 PyObject *a1, *a2; long r; \
Raymond Hettingerea3fdf42002-12-29 16:33:45 +000047 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
Guido van Rossum037b9401996-07-30 16:55:54 +000048 if(-1 == (r=AOP(a1,a2))) return NULL; \
Christian Heimes217cfd12007-12-02 14:31:20 +000049 return PyLong_FromLong(r); }
Guido van Rossum037b9401996-07-30 16:55:54 +000050
Thomas Wouters477c8d52006-05-27 19:21:47 +000051#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; \
Christian Heimes217cfd12007-12-02 14:31:20 +000055 return PyLong_FromSsize_t(r); }
Thomas Wouters477c8d52006-05-27 19:21:47 +000056
Guido van Rossum77f6a652002-04-03 22:41:51 +000057#define spami2b(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
58 PyObject *a1, *a2; long r; \
Raymond Hettingerea3fdf42002-12-29 16:33:45 +000059 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
Guido van Rossum77f6a652002-04-03 22:41:51 +000060 if(-1 == (r=AOP(a1,a2))) return NULL; \
61 return PyBool_FromLong(r); }
62
Fred Drake428e75f2001-08-09 20:14:34 +000063#define spamrc(OP,A) static PyObject *OP(PyObject *s, PyObject *a) { \
64 PyObject *a1, *a2; \
Raymond Hettingerea3fdf42002-12-29 16:33:45 +000065 if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
Fred Drake428e75f2001-08-09 20:14:34 +000066 return PyObject_RichCompare(a1,a2,A); }
67
Guido van Rossum037b9401996-07-30 16:55:54 +000068spami(truth , PyObject_IsTrue)
Guido van Rossumc9fb47e1996-08-21 17:40:51 +000069spam2(op_add , PyNumber_Add)
70spam2(op_sub , PyNumber_Subtract)
71spam2(op_mul , PyNumber_Multiply)
Fred Drake428e75f2001-08-09 20:14:34 +000072spam2(op_floordiv , PyNumber_FloorDivide)
73spam2(op_truediv , PyNumber_TrueDivide)
Guido van Rossumc9fb47e1996-08-21 17:40:51 +000074spam2(op_mod , PyNumber_Remainder)
75spam1(op_neg , PyNumber_Negative)
76spam1(op_pos , PyNumber_Positive)
77spam1(op_abs , PyNumber_Absolute)
78spam1(op_inv , PyNumber_Invert)
Fred Drakeea4d3f02000-09-17 16:09:27 +000079spam1(op_invert , PyNumber_Invert)
Guido van Rossumc9fb47e1996-08-21 17:40:51 +000080spam2(op_lshift , PyNumber_Lshift)
81spam2(op_rshift , PyNumber_Rshift)
Guido van Rossum99c185e1998-04-09 17:54:26 +000082spami(op_not_ , PyObject_Not)
Guido van Rossumc9fb47e1996-08-21 17:40:51 +000083spam2(op_and_ , PyNumber_And)
84spam2(op_xor , PyNumber_Xor)
85spam2(op_or_ , PyNumber_Or)
Armin Rigof5bd3b42005-12-29 16:50:42 +000086spam2(op_iadd , PyNumber_InPlaceAdd)
87spam2(op_isub , PyNumber_InPlaceSubtract)
88spam2(op_imul , PyNumber_InPlaceMultiply)
Armin Rigof5bd3b42005-12-29 16:50:42 +000089spam2(op_ifloordiv , PyNumber_InPlaceFloorDivide)
90spam2(op_itruediv , PyNumber_InPlaceTrueDivide)
91spam2(op_imod , PyNumber_InPlaceRemainder)
92spam2(op_ilshift , PyNumber_InPlaceLshift)
93spam2(op_irshift , PyNumber_InPlaceRshift)
94spam2(op_iand , PyNumber_InPlaceAnd)
95spam2(op_ixor , PyNumber_InPlaceXor)
96spam2(op_ior , PyNumber_InPlaceOr)
Guido van Rossumc9fb47e1996-08-21 17:40:51 +000097spam2(op_concat , PySequence_Concat)
Armin Rigof5bd3b42005-12-29 16:50:42 +000098spam2(op_iconcat , PySequence_InPlaceConcat)
Guido van Rossum77f6a652002-04-03 22:41:51 +000099spami2b(op_contains , PySequence_Contains)
Thomas Wouters477c8d52006-05-27 19:21:47 +0000100spamn2(indexOf , PySequence_Index)
101spamn2(countOf , PySequence_Count)
Guido van Rossumc9fb47e1996-08-21 17:40:51 +0000102spam2(op_getitem , PyObject_GetItem)
103spam2n(op_delitem , PyObject_DelItem)
104spam3n(op_setitem , PyObject_SetItem)
Fred Drake428e75f2001-08-09 20:14:34 +0000105spamrc(op_lt , Py_LT)
106spamrc(op_le , Py_LE)
107spamrc(op_eq , Py_EQ)
108spamrc(op_ne , Py_NE)
109spamrc(op_gt , Py_GT)
110spamrc(op_ge , Py_GE)
Guido van Rossum037b9401996-07-30 16:55:54 +0000111
112static PyObject*
Raymond Hettinger5959c552002-08-19 03:19:09 +0000113op_pow(PyObject *s, PyObject *a)
114{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000115 PyObject *a1, *a2;
116 if (PyArg_UnpackTuple(a,"pow", 2, 2, &a1, &a2))
117 return PyNumber_Power(a1, a2, Py_None);
118 return NULL;
Raymond Hettinger5959c552002-08-19 03:19:09 +0000119}
120
121static PyObject*
Armin Rigof5bd3b42005-12-29 16:50:42 +0000122op_ipow(PyObject *s, PyObject *a)
123{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000124 PyObject *a1, *a2;
125 if (PyArg_UnpackTuple(a,"ipow", 2, 2, &a1, &a2))
126 return PyNumber_InPlacePower(a1, a2, Py_None);
127 return NULL;
Armin Rigof5bd3b42005-12-29 16:50:42 +0000128}
129
Guido van Rossum38fff8c2006-03-07 18:50:55 +0000130static PyObject *
131op_index(PyObject *s, PyObject *a)
132{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000133 return PyNumber_Index(a);
Guido van Rossum38fff8c2006-03-07 18:50:55 +0000134}
135
Armin Rigof5bd3b42005-12-29 16:50:42 +0000136static PyObject*
Raymond Hettinger9543b342003-01-18 23:22:20 +0000137is_(PyObject *s, PyObject *a)
138{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000139 PyObject *a1, *a2, *result = NULL;
140 if (PyArg_UnpackTuple(a,"is_", 2, 2, &a1, &a2)) {
141 result = (a1 == a2) ? Py_True : Py_False;
142 Py_INCREF(result);
143 }
144 return result;
Raymond Hettinger9543b342003-01-18 23:22:20 +0000145}
146
147static PyObject*
148is_not(PyObject *s, PyObject *a)
149{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000150 PyObject *a1, *a2, *result = NULL;
151 if (PyArg_UnpackTuple(a,"is_not", 2, 2, &a1, &a2)) {
152 result = (a1 != a2) ? Py_True : Py_False;
153 Py_INCREF(result);
154 }
155 return result;
Raymond Hettinger9543b342003-01-18 23:22:20 +0000156}
157
Guido van Rossum037b9401996-07-30 16:55:54 +0000158#undef spam1
159#undef spam2
Raymond Hettinger36cd2bf2003-01-03 08:24:58 +0000160#undef spam1o
161#undef spam1o
Neal Norwitz200788c2002-08-13 22:20:41 +0000162#define spam1(OP,DOC) {#OP, OP, METH_VARARGS, PyDoc_STR(DOC)},
Armin Rigoc4308d52005-12-29 14:39:28 +0000163#define spam2(OP,ALTOP,DOC) {#OP, op_##OP, METH_VARARGS, PyDoc_STR(DOC)}, \
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000164 {#ALTOP, op_##OP, METH_VARARGS, PyDoc_STR(DOC)},
Raymond Hettinger36cd2bf2003-01-03 08:24:58 +0000165#define spam1o(OP,DOC) {#OP, OP, METH_O, PyDoc_STR(DOC)},
Armin Rigoc4308d52005-12-29 14:39:28 +0000166#define spam2o(OP,ALTOP,DOC) {#OP, op_##OP, METH_O, PyDoc_STR(DOC)}, \
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000167 {#ALTOP, op_##OP, METH_O, PyDoc_STR(DOC)},
Guido van Rossum037b9401996-07-30 16:55:54 +0000168
169static struct PyMethodDef operator_methods[] = {
Guido van Rossum037b9401996-07-30 16:55:54 +0000170
Raymond Hettinger36cd2bf2003-01-03 08:24:58 +0000171spam1o(truth,
Guido van Rossum77f6a652002-04-03 22:41:51 +0000172 "truth(a) -- Return True if a is true, False otherwise.")
Fred Drakeea4d3f02000-09-17 16:09:27 +0000173spam2(contains,__contains__,
Guido van Rossum832f6d21998-05-22 18:12:59 +0000174 "contains(a, b) -- Same as b in a (note reversed operands).")
Guido van Rossum17202301996-08-19 22:01:39 +0000175spam1(indexOf,
Guido van Rossum832f6d21998-05-22 18:12:59 +0000176 "indexOf(a, b) -- Return the first index of b in a.")
Guido van Rossum17202301996-08-19 22:01:39 +0000177spam1(countOf,
178 "countOf(a, b) -- Return the number of times b occurs in a.")
Guido van Rossum037b9401996-07-30 16:55:54 +0000179
Raymond Hettinger9543b342003-01-18 23:22:20 +0000180spam1(is_, "is_(a, b) -- Same as a is b.")
181spam1(is_not, "is_not(a, b) -- Same as a is not b.")
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000182spam2o(index, __index__, "index(a) -- Same as a.__index__()")
Guido van Rossum832f6d21998-05-22 18:12:59 +0000183spam2(add,__add__, "add(a, b) -- Same as a + b.")
184spam2(sub,__sub__, "sub(a, b) -- Same as a - b.")
185spam2(mul,__mul__, "mul(a, b) -- Same as a * b.")
Fred Drake428e75f2001-08-09 20:14:34 +0000186spam2(floordiv,__floordiv__, "floordiv(a, b) -- Same as a // b.")
Neal Norwitzbcc0db82006-03-24 08:14:36 +0000187spam2(truediv,__truediv__, "truediv(a, b) -- Same as a / b.")
Guido van Rossum832f6d21998-05-22 18:12:59 +0000188spam2(mod,__mod__, "mod(a, b) -- Same as a % b.")
Raymond Hettinger36cd2bf2003-01-03 08:24:58 +0000189spam2o(neg,__neg__, "neg(a) -- Same as -a.")
190spam2o(pos,__pos__, "pos(a) -- Same as +a.")
191spam2o(abs,__abs__, "abs(a) -- Same as abs(a).")
192spam2o(inv,__inv__, "inv(a) -- Same as ~a.")
193spam2o(invert,__invert__, "invert(a) -- Same as ~a.")
Guido van Rossum832f6d21998-05-22 18:12:59 +0000194spam2(lshift,__lshift__, "lshift(a, b) -- Same as a << b.")
195spam2(rshift,__rshift__, "rshift(a, b) -- Same as a >> b.")
Raymond Hettinger36cd2bf2003-01-03 08:24:58 +0000196spam2o(not_,__not__, "not_(a) -- Same as not a.")
Guido van Rossum832f6d21998-05-22 18:12:59 +0000197spam2(and_,__and__, "and_(a, b) -- Same as a & b.")
198spam2(xor,__xor__, "xor(a, b) -- Same as a ^ b.")
199spam2(or_,__or__, "or_(a, b) -- Same as a | b.")
Benjamin Petersona0dfa822009-11-13 02:25:08 +0000200spam2(iadd,__iadd__, "a = iadd(a, b) -- Same as a += b.")
201spam2(isub,__isub__, "a = isub(a, b) -- Same as a -= b.")
202spam2(imul,__imul__, "a = imul(a, b) -- Same as a *= b.")
203spam2(ifloordiv,__ifloordiv__, "a = ifloordiv(a, b) -- Same as a //= b.")
204spam2(itruediv,__itruediv__, "a = itruediv(a, b) -- Same as a /= b")
205spam2(imod,__imod__, "a = imod(a, b) -- Same as a %= b.")
206spam2(ilshift,__ilshift__, "a = ilshift(a, b) -- Same as a <<= b.")
207spam2(irshift,__irshift__, "a = irshift(a, b) -- Same as a >>= b.")
208spam2(iand,__iand__, "a = iand(a, b) -- Same as a &= b.")
209spam2(ixor,__ixor__, "a = ixor(a, b) -- Same as a ^= b.")
210spam2(ior,__ior__, "a = ior(a, b) -- Same as a |= b.")
Guido van Rossum17202301996-08-19 22:01:39 +0000211spam2(concat,__concat__,
Guido van Rossum832f6d21998-05-22 18:12:59 +0000212 "concat(a, b) -- Same as a + b, for a and b sequences.")
Armin Rigof5bd3b42005-12-29 16:50:42 +0000213spam2(iconcat,__iconcat__,
Benjamin Petersona0dfa822009-11-13 02:25:08 +0000214 "a = iconcat(a, b) -- Same as a += b, for a and b sequences.")
Guido van Rossum17202301996-08-19 22:01:39 +0000215spam2(getitem,__getitem__,
Guido van Rossum832f6d21998-05-22 18:12:59 +0000216 "getitem(a, b) -- Same as a[b].")
Guido van Rossum17202301996-08-19 22:01:39 +0000217spam2(setitem,__setitem__,
Guido van Rossum832f6d21998-05-22 18:12:59 +0000218 "setitem(a, b, c) -- Same as a[b] = c.")
Guido van Rossumc9fb47e1996-08-21 17:40:51 +0000219spam2(delitem,__delitem__,
Guido van Rossum832f6d21998-05-22 18:12:59 +0000220 "delitem(a, b) -- Same as del a[b].")
Armin Rigof5bd3b42005-12-29 16:50:42 +0000221spam2(pow,__pow__, "pow(a, b) -- Same as a ** b.")
Benjamin Petersona0dfa822009-11-13 02:25:08 +0000222spam2(ipow,__ipow__, "a = ipow(a, b) -- Same as a **= b.")
Fred Drake428e75f2001-08-09 20:14:34 +0000223spam2(lt,__lt__, "lt(a, b) -- Same as a<b.")
224spam2(le,__le__, "le(a, b) -- Same as a<=b.")
225spam2(eq,__eq__, "eq(a, b) -- Same as a==b.")
226spam2(ne,__ne__, "ne(a, b) -- Same as a!=b.")
227spam2(gt,__gt__, "gt(a, b) -- Same as a>b.")
228spam2(ge,__ge__, "ge(a, b) -- Same as a>=b.")
Guido van Rossum037b9401996-07-30 16:55:54 +0000229
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000230 {NULL, NULL} /* sentinel */
Guido van Rossum037b9401996-07-30 16:55:54 +0000231
Guido van Rossum037b9401996-07-30 16:55:54 +0000232};
233
Raymond Hettinger166958b2003-12-01 13:18:39 +0000234/* itemgetter object **********************************************************/
Guido van Rossum037b9401996-07-30 16:55:54 +0000235
Raymond Hettinger166958b2003-12-01 13:18:39 +0000236typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000237 PyObject_HEAD
238 Py_ssize_t nitems;
239 PyObject *item;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000240} itemgetterobject;
241
242static PyTypeObject itemgetter_type;
243
244static PyObject *
245itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
246{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000247 itemgetterobject *ig;
248 PyObject *item;
249 Py_ssize_t nitems;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000250
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000251 if (!_PyArg_NoKeywords("itemgetter()", kwds))
252 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000253
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000254 nitems = PyTuple_GET_SIZE(args);
255 if (nitems <= 1) {
256 if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &item))
257 return NULL;
258 } else
259 item = args;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000260
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000261 /* create itemgetterobject structure */
262 ig = PyObject_GC_New(itemgetterobject, &itemgetter_type);
263 if (ig == NULL)
264 return NULL;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000265
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000266 Py_INCREF(item);
267 ig->item = item;
268 ig->nitems = nitems;
269
270 PyObject_GC_Track(ig);
271 return (PyObject *)ig;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000272}
273
274static void
275itemgetter_dealloc(itemgetterobject *ig)
276{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000277 PyObject_GC_UnTrack(ig);
278 Py_XDECREF(ig->item);
279 PyObject_GC_Del(ig);
Raymond Hettinger166958b2003-12-01 13:18:39 +0000280}
281
282static int
283itemgetter_traverse(itemgetterobject *ig, visitproc visit, void *arg)
284{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000285 Py_VISIT(ig->item);
286 return 0;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000287}
288
289static PyObject *
290itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw)
291{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000292 PyObject *obj, *result;
293 Py_ssize_t i, nitems=ig->nitems;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj))
296 return NULL;
297 if (nitems == 1)
298 return PyObject_GetItem(obj, ig->item);
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 assert(PyTuple_Check(ig->item));
301 assert(PyTuple_GET_SIZE(ig->item) == nitems);
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000302
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000303 result = PyTuple_New(nitems);
304 if (result == NULL)
305 return NULL;
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000306
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000307 for (i=0 ; i < nitems ; i++) {
308 PyObject *item, *val;
309 item = PyTuple_GET_ITEM(ig->item, i);
310 val = PyObject_GetItem(obj, item);
311 if (val == NULL) {
312 Py_DECREF(result);
313 return NULL;
314 }
315 PyTuple_SET_ITEM(result, i, val);
316 }
317 return result;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000318}
319
320PyDoc_STRVAR(itemgetter_doc,
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000321"itemgetter(item, ...) --> itemgetter object\n\
Raymond Hettinger166958b2003-12-01 13:18:39 +0000322\n\
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000323Return a callable object that fetches the given item(s) from its operand.\n\
324After, f=itemgetter(2), the call f(r) returns r[2].\n\
325After, g=itemgetter(2,5,3), the call g(r) returns (r[2], r[5], r[3])");
Raymond Hettinger166958b2003-12-01 13:18:39 +0000326
327static PyTypeObject itemgetter_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000328 PyVarObject_HEAD_INIT(NULL, 0)
329 "operator.itemgetter", /* tp_name */
330 sizeof(itemgetterobject), /* tp_basicsize */
331 0, /* tp_itemsize */
332 /* methods */
333 (destructor)itemgetter_dealloc, /* tp_dealloc */
334 0, /* tp_print */
335 0, /* tp_getattr */
336 0, /* tp_setattr */
337 0, /* tp_reserved */
338 0, /* tp_repr */
339 0, /* tp_as_number */
340 0, /* tp_as_sequence */
341 0, /* tp_as_mapping */
342 0, /* tp_hash */
343 (ternaryfunc)itemgetter_call, /* tp_call */
344 0, /* tp_str */
345 PyObject_GenericGetAttr, /* tp_getattro */
346 0, /* tp_setattro */
347 0, /* tp_as_buffer */
348 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
349 itemgetter_doc, /* tp_doc */
350 (traverseproc)itemgetter_traverse, /* tp_traverse */
351 0, /* tp_clear */
352 0, /* tp_richcompare */
353 0, /* tp_weaklistoffset */
354 0, /* tp_iter */
355 0, /* tp_iternext */
356 0, /* tp_methods */
357 0, /* tp_members */
358 0, /* tp_getset */
359 0, /* tp_base */
360 0, /* tp_dict */
361 0, /* tp_descr_get */
362 0, /* tp_descr_set */
363 0, /* tp_dictoffset */
364 0, /* tp_init */
365 0, /* tp_alloc */
366 itemgetter_new, /* tp_new */
367 0, /* tp_free */
Raymond Hettinger166958b2003-12-01 13:18:39 +0000368};
369
370
371/* attrgetter object **********************************************************/
372
373typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000374 PyObject_HEAD
375 Py_ssize_t nattrs;
376 PyObject *attr;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000377} attrgetterobject;
378
379static PyTypeObject attrgetter_type;
380
381static PyObject *
382attrgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
383{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000384 attrgetterobject *ag;
385 PyObject *attr;
386 Py_ssize_t nattrs;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000387
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000388 if (!_PyArg_NoKeywords("attrgetter()", kwds))
389 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000390
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000391 nattrs = PyTuple_GET_SIZE(args);
392 if (nattrs <= 1) {
393 if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &attr))
394 return NULL;
395 } else
396 attr = args;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000397
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000398 /* create attrgetterobject structure */
399 ag = PyObject_GC_New(attrgetterobject, &attrgetter_type);
400 if (ag == NULL)
401 return NULL;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000402
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000403 Py_INCREF(attr);
404 ag->attr = attr;
405 ag->nattrs = nattrs;
406
407 PyObject_GC_Track(ag);
408 return (PyObject *)ag;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000409}
410
411static void
412attrgetter_dealloc(attrgetterobject *ag)
413{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000414 PyObject_GC_UnTrack(ag);
415 Py_XDECREF(ag->attr);
416 PyObject_GC_Del(ag);
Raymond Hettinger166958b2003-12-01 13:18:39 +0000417}
418
419static int
420attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg)
421{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000422 Py_VISIT(ag->attr);
423 return 0;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000424}
425
426static PyObject *
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000427dotted_getattr(PyObject *obj, PyObject *attr)
428{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000429 char *s, *p;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000430
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000431 if (!PyUnicode_Check(attr)) {
432 PyErr_SetString(PyExc_TypeError,
433 "attribute name must be a string");
434 return NULL;
435 }
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000437 s = _PyUnicode_AsString(attr);
438 Py_INCREF(obj);
439 for (;;) {
440 PyObject *newobj, *str;
441 p = strchr(s, '.');
442 str = p ? PyUnicode_FromStringAndSize(s, (p-s)) :
443 PyUnicode_FromString(s);
444 if (str == NULL) {
445 Py_DECREF(obj);
446 return NULL;
447 }
448 newobj = PyObject_GetAttr(obj, str);
449 Py_DECREF(str);
450 Py_DECREF(obj);
451 if (newobj == NULL)
452 return NULL;
453 obj = newobj;
454 if (p == NULL) break;
455 s = p+1;
456 }
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000457
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000458 return obj;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000459}
460
461static PyObject *
Raymond Hettinger166958b2003-12-01 13:18:39 +0000462attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
463{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000464 PyObject *obj, *result;
465 Py_ssize_t i, nattrs=ag->nattrs;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000467 if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj))
468 return NULL;
469 if (ag->nattrs == 1)
470 return dotted_getattr(obj, ag->attr);
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000471
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000472 assert(PyTuple_Check(ag->attr));
473 assert(PyTuple_GET_SIZE(ag->attr) == nattrs);
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000474
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000475 result = PyTuple_New(nattrs);
476 if (result == NULL)
477 return NULL;
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000478
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000479 for (i=0 ; i < nattrs ; i++) {
480 PyObject *attr, *val;
481 attr = PyTuple_GET_ITEM(ag->attr, i);
482 val = dotted_getattr(obj, attr);
483 if (val == NULL) {
484 Py_DECREF(result);
485 return NULL;
486 }
487 PyTuple_SET_ITEM(result, i, val);
488 }
489 return result;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000490}
491
492PyDoc_STRVAR(attrgetter_doc,
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000493"attrgetter(attr, ...) --> attrgetter object\n\
Raymond Hettinger166958b2003-12-01 13:18:39 +0000494\n\
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000495Return a callable object that fetches the given attribute(s) from its operand.\n\
496After, f=attrgetter('name'), the call f(r) returns r.name.\n\
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000497After, g=attrgetter('name', 'date'), the call g(r) returns (r.name, r.date).\n\
498After, h=attrgetter('name.first', 'name.last'), the call h(r) returns\n\
499(r.name.first, r.name.last).");
Raymond Hettinger166958b2003-12-01 13:18:39 +0000500
501static PyTypeObject attrgetter_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000502 PyVarObject_HEAD_INIT(NULL, 0)
503 "operator.attrgetter", /* tp_name */
504 sizeof(attrgetterobject), /* tp_basicsize */
505 0, /* tp_itemsize */
506 /* methods */
507 (destructor)attrgetter_dealloc, /* tp_dealloc */
508 0, /* tp_print */
509 0, /* tp_getattr */
510 0, /* tp_setattr */
511 0, /* tp_reserved */
512 0, /* tp_repr */
513 0, /* tp_as_number */
514 0, /* tp_as_sequence */
515 0, /* tp_as_mapping */
516 0, /* tp_hash */
517 (ternaryfunc)attrgetter_call, /* tp_call */
518 0, /* tp_str */
519 PyObject_GenericGetAttr, /* tp_getattro */
520 0, /* tp_setattro */
521 0, /* tp_as_buffer */
522 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
523 attrgetter_doc, /* tp_doc */
524 (traverseproc)attrgetter_traverse, /* tp_traverse */
525 0, /* tp_clear */
526 0, /* tp_richcompare */
527 0, /* tp_weaklistoffset */
528 0, /* tp_iter */
529 0, /* tp_iternext */
530 0, /* tp_methods */
531 0, /* tp_members */
532 0, /* tp_getset */
533 0, /* tp_base */
534 0, /* tp_dict */
535 0, /* tp_descr_get */
536 0, /* tp_descr_set */
537 0, /* tp_dictoffset */
538 0, /* tp_init */
539 0, /* tp_alloc */
540 attrgetter_new, /* tp_new */
541 0, /* tp_free */
Raymond Hettinger166958b2003-12-01 13:18:39 +0000542};
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000543
544
545/* methodcaller object **********************************************************/
546
547typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000548 PyObject_HEAD
549 PyObject *name;
550 PyObject *args;
551 PyObject *kwds;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000552} methodcallerobject;
553
554static PyTypeObject methodcaller_type;
555
556static PyObject *
557methodcaller_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
558{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000559 methodcallerobject *mc;
560 PyObject *name, *newargs;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000561
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000562 if (PyTuple_GET_SIZE(args) < 1) {
563 PyErr_SetString(PyExc_TypeError, "methodcaller needs at least "
564 "one argument, the method name");
565 return NULL;
566 }
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000567
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000568 /* create methodcallerobject structure */
569 mc = PyObject_GC_New(methodcallerobject, &methodcaller_type);
570 if (mc == NULL)
571 return NULL;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000572
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000573 newargs = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
574 if (newargs == NULL) {
575 Py_DECREF(mc);
576 return NULL;
577 }
578 mc->args = newargs;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000579
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000580 name = PyTuple_GET_ITEM(args, 0);
581 Py_INCREF(name);
582 mc->name = name;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000583
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000584 Py_XINCREF(kwds);
585 mc->kwds = kwds;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000586
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000587 PyObject_GC_Track(mc);
588 return (PyObject *)mc;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000589}
590
591static void
592methodcaller_dealloc(methodcallerobject *mc)
593{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000594 PyObject_GC_UnTrack(mc);
595 Py_XDECREF(mc->name);
596 Py_XDECREF(mc->args);
597 Py_XDECREF(mc->kwds);
598 PyObject_GC_Del(mc);
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000599}
600
601static int
602methodcaller_traverse(methodcallerobject *mc, visitproc visit, void *arg)
603{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000604 Py_VISIT(mc->args);
605 Py_VISIT(mc->kwds);
606 return 0;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000607}
608
609static PyObject *
610methodcaller_call(methodcallerobject *mc, PyObject *args, PyObject *kw)
611{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000612 PyObject *method, *obj, *result;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000613
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000614 if (!PyArg_UnpackTuple(args, "methodcaller", 1, 1, &obj))
615 return NULL;
616 method = PyObject_GetAttr(obj, mc->name);
617 if (method == NULL)
618 return NULL;
619 result = PyObject_Call(method, mc->args, mc->kwds);
620 Py_DECREF(method);
621 return result;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000622}
623
624PyDoc_STRVAR(methodcaller_doc,
625"methodcaller(name, ...) --> methodcaller object\n\
626\n\
627Return a callable object that calls the given method on its operand.\n\
628After, f = methodcaller('name'), the call f(r) returns r.name().\n\
629After, g = methodcaller('name', 'date', foo=1), the call g(r) returns\n\
630r.name('date', foo=1).");
631
632static PyTypeObject methodcaller_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 PyVarObject_HEAD_INIT(NULL, 0)
634 "operator.methodcaller", /* tp_name */
635 sizeof(methodcallerobject), /* tp_basicsize */
636 0, /* tp_itemsize */
637 /* methods */
638 (destructor)methodcaller_dealloc, /* tp_dealloc */
639 0, /* tp_print */
640 0, /* tp_getattr */
641 0, /* tp_setattr */
642 0, /* tp_reserved */
643 0, /* tp_repr */
644 0, /* tp_as_number */
645 0, /* tp_as_sequence */
646 0, /* tp_as_mapping */
647 0, /* tp_hash */
648 (ternaryfunc)methodcaller_call, /* tp_call */
649 0, /* tp_str */
650 PyObject_GenericGetAttr, /* tp_getattro */
651 0, /* tp_setattro */
652 0, /* tp_as_buffer */
653 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
654 methodcaller_doc, /* tp_doc */
655 (traverseproc)methodcaller_traverse, /* tp_traverse */
656 0, /* tp_clear */
657 0, /* tp_richcompare */
658 0, /* tp_weaklistoffset */
659 0, /* tp_iter */
660 0, /* tp_iternext */
661 0, /* tp_methods */
662 0, /* tp_members */
663 0, /* tp_getset */
664 0, /* tp_base */
665 0, /* tp_dict */
666 0, /* tp_descr_get */
667 0, /* tp_descr_set */
668 0, /* tp_dictoffset */
669 0, /* tp_init */
670 0, /* tp_alloc */
671 methodcaller_new, /* tp_new */
672 0, /* tp_free */
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000673};
674
675
Martin v. Löwis1a214512008-06-11 05:26:20 +0000676/* Initialization function for the module (*must* be called PyInit_operator) */
677
678
679static struct PyModuleDef operatormodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000680 PyModuleDef_HEAD_INIT,
681 "operator",
682 operator_doc,
683 -1,
684 operator_methods,
685 NULL,
686 NULL,
687 NULL,
688 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +0000689};
Guido van Rossum037b9401996-07-30 16:55:54 +0000690
Mark Hammondfe51c6d2002-08-02 02:27:13 +0000691PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +0000692PyInit_operator(void)
Guido van Rossum037b9401996-07-30 16:55:54 +0000693{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000694 PyObject *m;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000696 /* Create the module and add the functions */
697 m = PyModule_Create(&operatormodule);
698 if (m == NULL)
699 return NULL;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000700
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000701 if (PyType_Ready(&itemgetter_type) < 0)
702 return NULL;
703 Py_INCREF(&itemgetter_type);
704 PyModule_AddObject(m, "itemgetter", (PyObject *)&itemgetter_type);
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000705
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000706 if (PyType_Ready(&attrgetter_type) < 0)
707 return NULL;
708 Py_INCREF(&attrgetter_type);
709 PyModule_AddObject(m, "attrgetter", (PyObject *)&attrgetter_type);
710
711 if (PyType_Ready(&methodcaller_type) < 0)
712 return NULL;
713 Py_INCREF(&methodcaller_type);
714 PyModule_AddObject(m, "methodcaller", (PyObject *)&methodcaller_type);
715 return m;
Guido van Rossum037b9401996-07-30 16:55:54 +0000716}