blob: 866ec3a761e4722ef7452357d3efc5bfea7a9424 [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;
Antoine Pitroue9745712010-10-31 15:26:04 +0000386 Py_ssize_t nattrs, idx, char_idx;
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;
Antoine Pitroue9745712010-10-31 15:26:04 +0000395 }
396
397 attr = PyTuple_New(nattrs);
398 if (attr == NULL)
399 return NULL;
400
401 /* prepare attr while checking args */
402 for (idx = 0; idx < nattrs; ++idx) {
403 PyObject *item = PyTuple_GET_ITEM(args, idx);
404 Py_ssize_t item_len;
405 Py_UNICODE *item_buffer;
406 int dot_count;
407
408 if (!PyUnicode_Check(item)) {
409 PyErr_SetString(PyExc_TypeError,
410 "attribute name must be a string");
411 Py_DECREF(attr);
412 return NULL;
413 }
414 item_len = PyUnicode_GET_SIZE(item);
415 item_buffer = PyUnicode_AS_UNICODE(item);
416
417 /* check whethere the string is dotted */
418 dot_count = 0;
419 for (char_idx = 0; char_idx < item_len; ++char_idx) {
420 if (item_buffer[char_idx] == (Py_UNICODE)'.')
421 ++dot_count;
422 }
423
424 if (dot_count == 0) {
425 Py_INCREF(item);
426 PyUnicode_InternInPlace(&item);
427 PyTuple_SET_ITEM(attr, idx, item);
428 } else { /* make it a tuple of non-dotted attrnames */
429 PyObject *attr_chain = PyTuple_New(dot_count + 1);
430 PyObject *attr_chain_item;
Antoine Pitrou87298c42010-10-31 21:03:01 +0000431 Py_ssize_t unibuff_from = 0;
432 Py_ssize_t unibuff_till = 0;
433 Py_ssize_t attr_chain_idx = 0;
Antoine Pitroue9745712010-10-31 15:26:04 +0000434
435 if (attr_chain == NULL) {
436 Py_DECREF(attr);
437 return NULL;
438 }
439
Antoine Pitroue9745712010-10-31 15:26:04 +0000440 for (; dot_count > 0; --dot_count) {
441 while (item_buffer[unibuff_till] != (Py_UNICODE)'.') {
442 ++unibuff_till;
443 }
444 attr_chain_item = PyUnicode_FromUnicode(
445 item_buffer + unibuff_from,
446 unibuff_till - unibuff_from);
447 if (attr_chain_item == NULL) {
448 Py_DECREF(attr_chain);
449 Py_DECREF(attr);
450 return NULL;
451 }
452 PyUnicode_InternInPlace(&attr_chain_item);
453 PyTuple_SET_ITEM(attr_chain, attr_chain_idx, attr_chain_item);
454 ++attr_chain_idx;
455 unibuff_till = unibuff_from = unibuff_till + 1;
456 }
457
458 /* now add the last dotless name */
459 attr_chain_item = PyUnicode_FromUnicode(
460 item_buffer + unibuff_from,
461 item_len - unibuff_from);
462 if (attr_chain_item == NULL) {
463 Py_DECREF(attr_chain);
464 Py_DECREF(attr);
465 return NULL;
466 }
467 PyUnicode_InternInPlace(&attr_chain_item);
468 PyTuple_SET_ITEM(attr_chain, attr_chain_idx, attr_chain_item);
469
470 PyTuple_SET_ITEM(attr, idx, attr_chain);
471 }
472 }
Raymond Hettinger166958b2003-12-01 13:18:39 +0000473
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000474 /* create attrgetterobject structure */
475 ag = PyObject_GC_New(attrgetterobject, &attrgetter_type);
Antoine Pitroue9745712010-10-31 15:26:04 +0000476 if (ag == NULL) {
477 Py_DECREF(attr);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000478 return NULL;
Antoine Pitroue9745712010-10-31 15:26:04 +0000479 }
Raymond Hettinger166958b2003-12-01 13:18:39 +0000480
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000481 ag->attr = attr;
482 ag->nattrs = nattrs;
483
484 PyObject_GC_Track(ag);
485 return (PyObject *)ag;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000486}
487
488static void
489attrgetter_dealloc(attrgetterobject *ag)
490{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000491 PyObject_GC_UnTrack(ag);
492 Py_XDECREF(ag->attr);
493 PyObject_GC_Del(ag);
Raymond Hettinger166958b2003-12-01 13:18:39 +0000494}
495
496static int
497attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg)
498{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000499 Py_VISIT(ag->attr);
500 return 0;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000501}
502
503static PyObject *
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000504dotted_getattr(PyObject *obj, PyObject *attr)
505{
Antoine Pitroue9745712010-10-31 15:26:04 +0000506 PyObject *newobj;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000507
Antoine Pitroue9745712010-10-31 15:26:04 +0000508 /* attr is either a tuple or instance of str.
509 Ensured by the setup code of attrgetter_new */
510 if (PyTuple_CheckExact(attr)) { /* chained getattr */
511 Py_ssize_t name_idx = 0, name_count;
512 PyObject *attr_name;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000513
Antoine Pitroue9745712010-10-31 15:26:04 +0000514 name_count = PyTuple_GET_SIZE(attr);
515 Py_INCREF(obj);
516 for (name_idx = 0; name_idx < name_count; ++name_idx) {
517 attr_name = PyTuple_GET_ITEM(attr, name_idx);
518 newobj = PyObject_GetAttr(obj, attr_name);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000519 Py_DECREF(obj);
Antoine Pitroue9745712010-10-31 15:26:04 +0000520 if (newobj == NULL) {
521 return NULL;
522 }
523 /* here */
524 obj = newobj;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000525 }
Antoine Pitroue9745712010-10-31 15:26:04 +0000526 } else { /* single getattr */
527 newobj = PyObject_GetAttr(obj, attr);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000528 if (newobj == NULL)
529 return NULL;
530 obj = newobj;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000531 }
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000532
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000533 return obj;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000534}
535
536static PyObject *
Raymond Hettinger166958b2003-12-01 13:18:39 +0000537attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
538{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000539 PyObject *obj, *result;
540 Py_ssize_t i, nattrs=ag->nattrs;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000541
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000542 if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj))
543 return NULL;
Antoine Pitroue9745712010-10-31 15:26:04 +0000544 if (ag->nattrs == 1) /* ag->attr is always a tuple */
545 return dotted_getattr(obj, PyTuple_GET_ITEM(ag->attr, 0));
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000546
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000547 assert(PyTuple_Check(ag->attr));
548 assert(PyTuple_GET_SIZE(ag->attr) == nattrs);
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000549
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000550 result = PyTuple_New(nattrs);
551 if (result == NULL)
552 return NULL;
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000553
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000554 for (i=0 ; i < nattrs ; i++) {
555 PyObject *attr, *val;
556 attr = PyTuple_GET_ITEM(ag->attr, i);
557 val = dotted_getattr(obj, attr);
558 if (val == NULL) {
559 Py_DECREF(result);
560 return NULL;
561 }
562 PyTuple_SET_ITEM(result, i, val);
563 }
564 return result;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000565}
566
567PyDoc_STRVAR(attrgetter_doc,
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000568"attrgetter(attr, ...) --> attrgetter object\n\
Raymond Hettinger166958b2003-12-01 13:18:39 +0000569\n\
Raymond Hettinger984f9bb2005-03-09 16:38:48 +0000570Return a callable object that fetches the given attribute(s) from its operand.\n\
571After, f=attrgetter('name'), the call f(r) returns r.name.\n\
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000572After, g=attrgetter('name', 'date'), the call g(r) returns (r.name, r.date).\n\
573After, h=attrgetter('name.first', 'name.last'), the call h(r) returns\n\
574(r.name.first, r.name.last).");
Raymond Hettinger166958b2003-12-01 13:18:39 +0000575
576static PyTypeObject attrgetter_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000577 PyVarObject_HEAD_INIT(NULL, 0)
578 "operator.attrgetter", /* tp_name */
579 sizeof(attrgetterobject), /* tp_basicsize */
580 0, /* tp_itemsize */
581 /* methods */
582 (destructor)attrgetter_dealloc, /* tp_dealloc */
583 0, /* tp_print */
584 0, /* tp_getattr */
585 0, /* tp_setattr */
586 0, /* tp_reserved */
587 0, /* tp_repr */
588 0, /* tp_as_number */
589 0, /* tp_as_sequence */
590 0, /* tp_as_mapping */
591 0, /* tp_hash */
592 (ternaryfunc)attrgetter_call, /* tp_call */
593 0, /* tp_str */
594 PyObject_GenericGetAttr, /* tp_getattro */
595 0, /* tp_setattro */
596 0, /* tp_as_buffer */
597 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
598 attrgetter_doc, /* tp_doc */
599 (traverseproc)attrgetter_traverse, /* tp_traverse */
600 0, /* tp_clear */
601 0, /* tp_richcompare */
602 0, /* tp_weaklistoffset */
603 0, /* tp_iter */
604 0, /* tp_iternext */
605 0, /* tp_methods */
606 0, /* tp_members */
607 0, /* tp_getset */
608 0, /* tp_base */
609 0, /* tp_dict */
610 0, /* tp_descr_get */
611 0, /* tp_descr_set */
612 0, /* tp_dictoffset */
613 0, /* tp_init */
614 0, /* tp_alloc */
615 attrgetter_new, /* tp_new */
616 0, /* tp_free */
Raymond Hettinger166958b2003-12-01 13:18:39 +0000617};
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000618
619
620/* methodcaller object **********************************************************/
621
622typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000623 PyObject_HEAD
624 PyObject *name;
625 PyObject *args;
626 PyObject *kwds;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000627} methodcallerobject;
628
629static PyTypeObject methodcaller_type;
630
631static PyObject *
632methodcaller_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
633{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000634 methodcallerobject *mc;
635 PyObject *name, *newargs;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000636
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000637 if (PyTuple_GET_SIZE(args) < 1) {
638 PyErr_SetString(PyExc_TypeError, "methodcaller needs at least "
639 "one argument, the method name");
640 return NULL;
641 }
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000643 /* create methodcallerobject structure */
644 mc = PyObject_GC_New(methodcallerobject, &methodcaller_type);
645 if (mc == NULL)
646 return NULL;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000647
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000648 newargs = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
649 if (newargs == NULL) {
650 Py_DECREF(mc);
651 return NULL;
652 }
653 mc->args = newargs;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000654
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000655 name = PyTuple_GET_ITEM(args, 0);
656 Py_INCREF(name);
657 mc->name = name;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000659 Py_XINCREF(kwds);
660 mc->kwds = kwds;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 PyObject_GC_Track(mc);
663 return (PyObject *)mc;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000664}
665
666static void
667methodcaller_dealloc(methodcallerobject *mc)
668{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000669 PyObject_GC_UnTrack(mc);
670 Py_XDECREF(mc->name);
671 Py_XDECREF(mc->args);
672 Py_XDECREF(mc->kwds);
673 PyObject_GC_Del(mc);
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000674}
675
676static int
677methodcaller_traverse(methodcallerobject *mc, visitproc visit, void *arg)
678{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 Py_VISIT(mc->args);
680 Py_VISIT(mc->kwds);
681 return 0;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000682}
683
684static PyObject *
685methodcaller_call(methodcallerobject *mc, PyObject *args, PyObject *kw)
686{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 PyObject *method, *obj, *result;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000688
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000689 if (!PyArg_UnpackTuple(args, "methodcaller", 1, 1, &obj))
690 return NULL;
691 method = PyObject_GetAttr(obj, mc->name);
692 if (method == NULL)
693 return NULL;
694 result = PyObject_Call(method, mc->args, mc->kwds);
695 Py_DECREF(method);
696 return result;
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000697}
698
699PyDoc_STRVAR(methodcaller_doc,
700"methodcaller(name, ...) --> methodcaller object\n\
701\n\
702Return a callable object that calls the given method on its operand.\n\
703After, f = methodcaller('name'), the call f(r) returns r.name().\n\
704After, g = methodcaller('name', 'date', foo=1), the call g(r) returns\n\
705r.name('date', foo=1).");
706
707static PyTypeObject methodcaller_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000708 PyVarObject_HEAD_INIT(NULL, 0)
709 "operator.methodcaller", /* tp_name */
710 sizeof(methodcallerobject), /* tp_basicsize */
711 0, /* tp_itemsize */
712 /* methods */
713 (destructor)methodcaller_dealloc, /* tp_dealloc */
714 0, /* tp_print */
715 0, /* tp_getattr */
716 0, /* tp_setattr */
717 0, /* tp_reserved */
718 0, /* tp_repr */
719 0, /* tp_as_number */
720 0, /* tp_as_sequence */
721 0, /* tp_as_mapping */
722 0, /* tp_hash */
723 (ternaryfunc)methodcaller_call, /* tp_call */
724 0, /* tp_str */
725 PyObject_GenericGetAttr, /* tp_getattro */
726 0, /* tp_setattro */
727 0, /* tp_as_buffer */
728 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
729 methodcaller_doc, /* tp_doc */
730 (traverseproc)methodcaller_traverse, /* tp_traverse */
731 0, /* tp_clear */
732 0, /* tp_richcompare */
733 0, /* tp_weaklistoffset */
734 0, /* tp_iter */
735 0, /* tp_iternext */
736 0, /* tp_methods */
737 0, /* tp_members */
738 0, /* tp_getset */
739 0, /* tp_base */
740 0, /* tp_dict */
741 0, /* tp_descr_get */
742 0, /* tp_descr_set */
743 0, /* tp_dictoffset */
744 0, /* tp_init */
745 0, /* tp_alloc */
746 methodcaller_new, /* tp_new */
747 0, /* tp_free */
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000748};
749
750
Martin v. Löwis1a214512008-06-11 05:26:20 +0000751/* Initialization function for the module (*must* be called PyInit_operator) */
752
753
754static struct PyModuleDef operatormodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000755 PyModuleDef_HEAD_INIT,
756 "operator",
757 operator_doc,
758 -1,
759 operator_methods,
760 NULL,
761 NULL,
762 NULL,
763 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +0000764};
Guido van Rossum037b9401996-07-30 16:55:54 +0000765
Mark Hammondfe51c6d2002-08-02 02:27:13 +0000766PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +0000767PyInit_operator(void)
Guido van Rossum037b9401996-07-30 16:55:54 +0000768{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 PyObject *m;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000771 /* Create the module and add the functions */
772 m = PyModule_Create(&operatormodule);
773 if (m == NULL)
774 return NULL;
Raymond Hettinger166958b2003-12-01 13:18:39 +0000775
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000776 if (PyType_Ready(&itemgetter_type) < 0)
777 return NULL;
778 Py_INCREF(&itemgetter_type);
779 PyModule_AddObject(m, "itemgetter", (PyObject *)&itemgetter_type);
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000780
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000781 if (PyType_Ready(&attrgetter_type) < 0)
782 return NULL;
783 Py_INCREF(&attrgetter_type);
784 PyModule_AddObject(m, "attrgetter", (PyObject *)&attrgetter_type);
785
786 if (PyType_Ready(&methodcaller_type) < 0)
787 return NULL;
788 Py_INCREF(&methodcaller_type);
789 PyModule_AddObject(m, "methodcaller", (PyObject *)&methodcaller_type);
790 return m;
Guido van Rossum037b9401996-07-30 16:55:54 +0000791}