Modules/operator.c

   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
  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
  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
  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*
 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 *
 159 op_index(PyObject *s, PyObject *a)
 160 {
 161         return PyNumber_Index(a);
 162 }
 163
 164 static PyObject*
 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*
 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*
 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*
 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*
 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 static struct PyMethodDef operator_methods[] = {
 239
 240 spam1o(isCallable,
 241  "isCallable(a) -- Same as callable(a).")
 242 spam1o(isNumberType,
 243  "isNumberType(a) -- Return True if a has a numeric type, False otherwise.")
 244 spam1o(isSequenceType,
 245  "isSequenceType(a) -- Return True if a has a sequence type, False otherwise.")
 246 spam1o(truth,
 247  "truth(a) -- Return True if a is true, False otherwise.")
 248 spam2(contains,__contains__,
 249  "contains(a, b) -- Same as b in a (note reversed operands).")
 250 spam1(sequenceIncludes,
 251  "sequenceIncludes(a, b) -- Same as b in a (note reversed operands; deprecated).")
 252 spam1(indexOf,
 253  "indexOf(a, b) -- Return the first index of b in a.")
 254 spam1(countOf,
 255  "countOf(a, b) -- Return the number of times b occurs in a.")
 256 spam1o(isMappingType,
 257  "isMappingType(a) -- Return True if a has a mapping type, False otherwise.")
 258
 259 spam1(is_, "is_(a, b) -- Same as a is b.")
 260 spam1(is_not, "is_not(a, b) -- Same as a is not b.")
 261 spam2o(index, __index__, "index(a) -- Same as a.__index__()")
 262 spam2(add,__add__, "add(a, b) -- Same as a + b.")
 263 spam2(sub,__sub__, "sub(a, b) -- Same as a - b.")
 264 spam2(mul,__mul__, "mul(a, b) -- Same as a * b.")
 265 spam2(div,__div__, "div(a, b) -- Same as a / b when __future__.division is not in effect.")
 266 spam2(floordiv,__floordiv__, "floordiv(a, b) -- Same as a // b.")
 267 spam2(truediv,__truediv__, "truediv(a, b) -- Same as a / b when __future__.division is in effect.")
 268 spam2(mod,__mod__, "mod(a, b) -- Same as a % b.")
 269 spam2o(neg,__neg__, "neg(a) -- Same as -a.")
 270 spam2o(pos,__pos__, "pos(a) -- Same as +a.")
 271 spam2o(abs,__abs__, "abs(a) -- Same as abs(a).")
 272 spam2o(inv,__inv__, "inv(a) -- Same as ~a.")
 273 spam2o(invert,__invert__, "invert(a) -- Same as ~a.")
 274 spam2(lshift,__lshift__, "lshift(a, b) -- Same as a << b.")
 275 spam2(rshift,__rshift__, "rshift(a, b) -- Same as a >> b.")
 276 spam2o(not_,__not__, "not_(a) -- Same as not a.")
 277 spam2(and_,__and__, "and_(a, b) -- Same as a & b.")
 278 spam2(xor,__xor__, "xor(a, b) -- Same as a ^ b.")
 279 spam2(or_,__or__, "or_(a, b) -- Same as a | b.")
 280 spam2(iadd,__iadd__, "a = iadd(a, b) -- Same as a += b.")
 281 spam2(isub,__isub__, "a = isub(a, b) -- Same as a -= b.")
 282 spam2(imul,__imul__, "a = imul(a, b) -- Same as a *= b.")
 283 spam2(idiv,__idiv__, "a = idiv(a, b) -- Same as a /= b when __future__.division is not in effect.")
 284 spam2(ifloordiv,__ifloordiv__, "a = ifloordiv(a, b) -- Same as a //= b.")
 285 spam2(itruediv,__itruediv__, "a = itruediv(a, b) -- Same as a /= b when __future__.division is in effect.")
 286 spam2(imod,__imod__, "a = imod(a, b) -- Same as a %= b.")
 287 spam2(ilshift,__ilshift__, "a = ilshift(a, b) -- Same as a <<= b.")
 288 spam2(irshift,__irshift__, "a = irshift(a, b) -- Same as a >>= b.")
 289 spam2(iand,__iand__, "a = iand(a, b) -- Same as a &= b.")
 290 spam2(ixor,__ixor__, "a = ixor(a, b) -- Same as a ^= b.")
 291 spam2(ior,__ior__, "a = ior(a, b) -- Same as a |= b.")
 292 spam2(concat,__concat__,
 293  "concat(a, b) -- Same as a + b, for a and b sequences.")
 294 spam2(repeat,__repeat__,
 295  "repeat(a, b) -- Return a * b, where a is a sequence, and b is an integer.")
 296 spam2(iconcat,__iconcat__,
 297  "a = iconcat(a, b) -- Same as a += b, for a and b sequences.")
 298 spam2(irepeat,__irepeat__,
 299  "a = irepeat(a, b) -- Same as a *= b, where a is a sequence, and b is an integer.")
 300 spam2(getitem,__getitem__,
 301  "getitem(a, b) -- Same as a[b].")
 302 spam2(setitem,__setitem__,
 303  "setitem(a, b, c) -- Same as a[b] = c.")
 304 spam2(delitem,__delitem__,
 305  "delitem(a, b) -- Same as del a[b].")
 306 spam2(pow,__pow__, "pow(a, b) -- Same as a ** b.")
 307 spam2(ipow,__ipow__, "a = ipow(a, b) -- Same as a **= b.")
 308 spam2(getslice,__getslice__,
 309  "getslice(a, b, c) -- Same as a[b:c].")
 310 spam2(setslice,__setslice__,
 311 "setslice(a, b, c, d) -- Same as a[b:c] = d.")
 312 spam2(delslice,__delslice__,
 313 "delslice(a, b, c) -- Same as del a[b:c].")
 314 spam2(lt,__lt__, "lt(a, b) -- Same as a<b.")
 315 spam2(le,__le__, "le(a, b) -- Same as a<=b.")
 316 spam2(eq,__eq__, "eq(a, b) -- Same as a==b.")
 317 spam2(ne,__ne__, "ne(a, b) -- Same as a!=b.")
 318 spam2(gt,__gt__, "gt(a, b) -- Same as a>b.")
 319 spam2(ge,__ge__, "ge(a, b) -- Same as a>=b.")
 320
 321         {NULL,          NULL}           /* sentinel */
 322
 323 };
 324
 325 /* itemgetter object **********************************************************/
 326
 327 typedef struct {
 328         PyObject_HEAD
 329         Py_ssize_t nitems;
 330         PyObject *item;
 331 } itemgetterobject;
 332
 333 static PyTypeObject itemgetter_type;
 334
 335 static PyObject *
 336 itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 337 {
 338         itemgetterobject *ig;
 339         PyObject *item;
 340         Py_ssize_t nitems;
 341
 342         if (!_PyArg_NoKeywords("itemgetter()", kwds))
 343                 return NULL;
 344
 345         nitems = PyTuple_GET_SIZE(args);
 346         if (nitems <= 1) {
 347                 if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &item))
 348                         return NULL;
 349         } else
 350                 item = args;
 351
 352         /* create itemgetterobject structure */
 353         ig = PyObject_GC_New(itemgetterobject, &itemgetter_type);
 354         if (ig == NULL)
 355                 return NULL;
 356
 357         Py_INCREF(item);
 358         ig->item = item;
 359         ig->nitems = nitems;
 360
 361         PyObject_GC_Track(ig);
 362         return (PyObject *)ig;
 363 }
 364
 365 static void
 366 itemgetter_dealloc(itemgetterobject *ig)
 367 {
 368         PyObject_GC_UnTrack(ig);
 369         Py_XDECREF(ig->item);
 370         PyObject_GC_Del(ig);
 371 }
 372
 373 static int
 374 itemgetter_traverse(itemgetterobject *ig, visitproc visit, void *arg)
 375 {
 376         Py_VISIT(ig->item);
 377         return 0;
 378 }
 379
 380 static PyObject *
 381 itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw)
 382 {
 383         PyObject *obj, *result;
 384         Py_ssize_t i, nitems=ig->nitems;
 385
 386         if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj))
 387                 return NULL;
 388         if (nitems == 1)
 389                 return PyObject_GetItem(obj, ig->item);
 390
 391         assert(PyTuple_Check(ig->item));
 392         assert(PyTuple_GET_SIZE(ig->item) == nitems);
 393
 394         result = PyTuple_New(nitems);
 395         if (result == NULL)
 396                 return NULL;
 397
 398         for (i=0 ; i < nitems ; i++) {
 399                 PyObject *item, *val;
 400                 item = PyTuple_GET_ITEM(ig->item, i);
 401                 val = PyObject_GetItem(obj, item);
 402                 if (val == NULL) {
 403                         Py_DECREF(result);
 404                         return NULL;
 405                 }
 406                 PyTuple_SET_ITEM(result, i, val);
 407         }
 408         return result;
 409 }
 410
 411 PyDoc_STRVAR(itemgetter_doc,
 412 "itemgetter(item, ...) --> itemgetter object\n\
 413 \n\
 414 Return a callable object that fetches the given item(s) from its operand.\n\
 415 After, f=itemgetter(2), the call f(r) returns r[2].\n\
 416 After, g=itemgetter(2,5,3), the call g(r) returns (r[2], r[5], r[3])");
 417
 418 static PyTypeObject itemgetter_type = {
 419         PyVarObject_HEAD_INIT(NULL, 0)
 420         "operator.itemgetter",          /* tp_name */
 421         sizeof(itemgetterobject),       /* tp_basicsize */
 422         0,                              /* tp_itemsize */
 423         /* methods */
 424         (destructor)itemgetter_dealloc, /* tp_dealloc */
 425         0,                              /* tp_print */
 426         0,                              /* tp_getattr */
 427         0,                              /* tp_setattr */
 428         0,                              /* tp_compare */
 429         0,                              /* tp_repr */
 430         0,                              /* tp_as_number */
 431         0,                              /* tp_as_sequence */
 432         0,                              /* tp_as_mapping */
 433         0,                              /* tp_hash */
 434         (ternaryfunc)itemgetter_call,   /* tp_call */
 435         0,                              /* tp_str */
 436         PyObject_GenericGetAttr,        /* tp_getattro */
 437         0,                              /* tp_setattro */
 438         0,                              /* tp_as_buffer */
 439         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,        /* tp_flags */
 440         itemgetter_doc,                 /* tp_doc */
 441         (traverseproc)itemgetter_traverse,      /* tp_traverse */
 442         0,                              /* tp_clear */
 443         0,                              /* tp_richcompare */
 444         0,                              /* tp_weaklistoffset */
 445         0,                              /* tp_iter */
 446         0,                              /* tp_iternext */
 447         0,                              /* tp_methods */
 448         0,                              /* tp_members */
 449         0,                              /* tp_getset */
 450         0,                              /* tp_base */
 451         0,                              /* tp_dict */
 452         0,                              /* tp_descr_get */
 453         0,                              /* tp_descr_set */
 454         0,                              /* tp_dictoffset */
 455         0,                              /* tp_init */
 456         0,                              /* tp_alloc */
 457         itemgetter_new,                 /* tp_new */
 458         0,                              /* tp_free */
 459 };
 460
 461
 462 /* attrgetter object **********************************************************/
 463
 464 typedef struct {
 465         PyObject_HEAD
 466         Py_ssize_t nattrs;
 467         PyObject *attr;
 468 } attrgetterobject;
 469
 470 static PyTypeObject attrgetter_type;
 471
 472 static PyObject *
 473 attrgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 474 {
 475         attrgetterobject *ag;
 476         PyObject *attr;
 477         Py_ssize_t nattrs;
 478
 479         if (!_PyArg_NoKeywords("attrgetter()", kwds))
 480                 return NULL;
 481
 482         nattrs = PyTuple_GET_SIZE(args);
 483         if (nattrs <= 1) {
 484                 if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &attr))
 485                         return NULL;
 486         } else
 487                 attr = args;
 488
 489         /* create attrgetterobject structure */
 490         ag = PyObject_GC_New(attrgetterobject, &attrgetter_type);
 491         if (ag == NULL)
 492                 return NULL;
 493
 494         Py_INCREF(attr);
 495         ag->attr = attr;
 496         ag->nattrs = nattrs;
 497
 498         PyObject_GC_Track(ag);
 499         return (PyObject *)ag;
 500 }
 501
 502 static void
 503 attrgetter_dealloc(attrgetterobject *ag)
 504 {
 505         PyObject_GC_UnTrack(ag);
 506         Py_XDECREF(ag->attr);
 507         PyObject_GC_Del(ag);
 508 }
 509
 510 static int
 511 attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg)
 512 {
 513         Py_VISIT(ag->attr);
 514         return 0;
 515 }
 516
 517 static PyObject *
 518 dotted_getattr(PyObject *obj, PyObject *attr)
 519 {
 520         char *s, *p;
 521
 522 #ifdef Py_USING_UNICODE
 523         if (PyUnicode_Check(attr)) {
 524                 attr = _PyUnicode_AsDefaultEncodedString(attr, NULL);
 525                 if (attr == NULL)
 526                         return NULL;
 527         }
 528 #endif
 529
 530         if (!PyString_Check(attr)) {
 531                 PyErr_SetString(PyExc_TypeError,
 532                                 "attribute name must be a string");
 533                 return NULL;
 534         }
 535
 536         s = PyString_AS_STRING(attr);
 537         Py_INCREF(obj);
 538         for (;;) {
 539                 PyObject *newobj, *str;
 540                 p = strchr(s, '.');
 541                 str = p ? PyString_FromStringAndSize(s, (p-s)) :
 542                           PyString_FromString(s);
 543                 if (str == NULL) {
 544                         Py_DECREF(obj);
 545                         return NULL;
 546                 }
 547                 newobj = PyObject_GetAttr(obj, str);
 548                 Py_DECREF(str);
 549                 Py_DECREF(obj);
 550                 if (newobj == NULL)
 551                         return NULL;
 552                 obj = newobj;
 553                 if (p == NULL) break;
 554                 s = p+1;
 555         }
 556
 557         return obj;
 558 }
 559
 560 static PyObject *
 561 attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
 562 {
 563         PyObject *obj, *result;
 564         Py_ssize_t i, nattrs=ag->nattrs;
 565
 566         if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj))
 567                 return NULL;
 568         if (ag->nattrs == 1)
 569                 return dotted_getattr(obj, ag->attr);
 570
 571         assert(PyTuple_Check(ag->attr));
 572         assert(PyTuple_GET_SIZE(ag->attr) == nattrs);
 573
 574         result = PyTuple_New(nattrs);
 575         if (result == NULL)
 576                 return NULL;
 577
 578         for (i=0 ; i < nattrs ; i++) {
 579                 PyObject *attr, *val;
 580                 attr = PyTuple_GET_ITEM(ag->attr, i);
 581                 val = dotted_getattr(obj, attr);
 582                 if (val == NULL) {
 583                         Py_DECREF(result);
 584                         return NULL;
 585                 }
 586                 PyTuple_SET_ITEM(result, i, val);
 587         }
 588         return result;
 589 }
 590
 591 PyDoc_STRVAR(attrgetter_doc,
 592 "attrgetter(attr, ...) --> attrgetter object\n\
 593 \n\
 594 Return a callable object that fetches the given attribute(s) from its operand.\n\
 595 After, f=attrgetter('name'), the call f(r) returns r.name.\n\
 596 After, g=attrgetter('name', 'date'), the call g(r) returns (r.name, r.date).\n\
 597 After, h=attrgetter('name.first', 'name.last'), the call h(r) returns\n\
 598 (r.name.first, r.name.last).");
 599
 600 static PyTypeObject attrgetter_type = {
 601         PyVarObject_HEAD_INIT(NULL, 0)
 602         "operator.attrgetter",          /* tp_name */
 603         sizeof(attrgetterobject),       /* tp_basicsize */
 604         0,                              /* tp_itemsize */
 605         /* methods */
 606         (destructor)attrgetter_dealloc, /* tp_dealloc */
 607         0,                              /* tp_print */
 608         0,                              /* tp_getattr */
 609         0,                              /* tp_setattr */
 610         0,                              /* tp_compare */
 611         0,                              /* tp_repr */
 612         0,                              /* tp_as_number */
 613         0,                              /* tp_as_sequence */
 614         0,                              /* tp_as_mapping */
 615         0,                              /* tp_hash */
 616         (ternaryfunc)attrgetter_call,   /* tp_call */
 617         0,                              /* tp_str */
 618         PyObject_GenericGetAttr,        /* tp_getattro */
 619         0,                              /* tp_setattro */
 620         0,                              /* tp_as_buffer */
 621         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,        /* tp_flags */
 622         attrgetter_doc,                 /* tp_doc */
 623         (traverseproc)attrgetter_traverse,      /* tp_traverse */
 624         0,                              /* tp_clear */
 625         0,                              /* tp_richcompare */
 626         0,                              /* tp_weaklistoffset */
 627         0,                              /* tp_iter */
 628         0,                              /* tp_iternext */
 629         0,                              /* tp_methods */
 630         0,                              /* tp_members */
 631         0,                              /* tp_getset */
 632         0,                              /* tp_base */
 633         0,                              /* tp_dict */
 634         0,                              /* tp_descr_get */
 635         0,                              /* tp_descr_set */
 636         0,                              /* tp_dictoffset */
 637         0,                              /* tp_init */
 638         0,                              /* tp_alloc */
 639         attrgetter_new,                 /* tp_new */
 640         0,                              /* tp_free */
 641 };
 642
 643
 644 /* methodcaller object **********************************************************/
 645
 646 typedef struct {
 647         PyObject_HEAD
 648         PyObject *name;
 649         PyObject *args;
 650         PyObject *kwds;
 651 } methodcallerobject;
 652
 653 static PyTypeObject methodcaller_type;
 654
 655 static PyObject *
 656 methodcaller_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 657 {
 658         methodcallerobject *mc;
 659         PyObject *name, *newargs;
 660
 661         if (PyTuple_GET_SIZE(args) < 1) {
 662                 PyErr_SetString(PyExc_TypeError, "methodcaller needs at least "
 663                                 "one argument, the method name");
 664                 return NULL;
 665         }
 666
 667         /* create methodcallerobject structure */
 668         mc = PyObject_GC_New(methodcallerobject, &methodcaller_type);
 669         if (mc == NULL)
 670                 return NULL;
 671
 672         newargs = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
 673         if (newargs == NULL) {
 674                 Py_DECREF(mc);
 675                 return NULL;
 676         }
 677         mc->args = newargs;
 678
 679         name = PyTuple_GET_ITEM(args, 0);
 680         Py_INCREF(name);
 681         mc->name = name;
 682
 683         Py_XINCREF(kwds);
 684         mc->kwds = kwds;
 685
 686         PyObject_GC_Track(mc);
 687         return (PyObject *)mc;
 688 }
 689
 690 static void
 691 methodcaller_dealloc(methodcallerobject *mc)
 692 {
 693         PyObject_GC_UnTrack(mc);
 694         Py_XDECREF(mc->name);
 695         Py_XDECREF(mc->args);
 696         Py_XDECREF(mc->kwds);
 697         PyObject_GC_Del(mc);
 698 }
 699
 700 static int
 701 methodcaller_traverse(methodcallerobject *mc, visitproc visit, void *arg)
 702 {
 703         Py_VISIT(mc->args);
 704         Py_VISIT(mc->kwds);
 705         return 0;
 706 }
 707
 708 static PyObject *
 709 methodcaller_call(methodcallerobject *mc, PyObject *args, PyObject *kw)
 710 {
 711         PyObject *method, *obj, *result;
 712
 713         if (!PyArg_UnpackTuple(args, "methodcaller", 1, 1, &obj))
 714                 return NULL;
 715         method = PyObject_GetAttr(obj, mc->name);
 716         if (method == NULL)
 717                 return NULL;
 718         result = PyObject_Call(method, mc->args, mc->kwds);
 719         Py_DECREF(method);
 720         return result;
 721 }
 722
 723 PyDoc_STRVAR(methodcaller_doc,
 724 "methodcaller(name, ...) --> methodcaller object\n\
 725 \n\
 726 Return a callable object that calls the given method on its operand.\n\
 727 After, f = methodcaller('name'), the call f(r) returns r.name().\n\
 728 After, g = methodcaller('name', 'date', foo=1), the call g(r) returns\n\
 729 r.name('date', foo=1).");
 730
 731 static PyTypeObject methodcaller_type = {
 732         PyVarObject_HEAD_INIT(NULL, 0)
 733         "operator.methodcaller",        /* tp_name */
 734         sizeof(methodcallerobject),     /* tp_basicsize */
 735         0,                              /* tp_itemsize */
 736         /* methods */
 737         (destructor)methodcaller_dealloc, /* tp_dealloc */
 738         0,                              /* tp_print */
 739         0,                              /* tp_getattr */
 740         0,                              /* tp_setattr */
 741         0,                              /* tp_compare */
 742         0,                              /* tp_repr */
 743         0,                              /* tp_as_number */
 744         0,                              /* tp_as_sequence */
 745         0,                              /* tp_as_mapping */
 746         0,                              /* tp_hash */
 747         (ternaryfunc)methodcaller_call, /* tp_call */
 748         0,                              /* tp_str */
 749         PyObject_GenericGetAttr,        /* tp_getattro */
 750         0,                              /* tp_setattro */
 751         0,                              /* tp_as_buffer */
 752         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
 753         methodcaller_doc,                       /* tp_doc */
 754         (traverseproc)methodcaller_traverse,    /* tp_traverse */
 755         0,                              /* tp_clear */
 756         0,                              /* tp_richcompare */
 757         0,                              /* tp_weaklistoffset */
 758         0,                              /* tp_iter */
 759         0,                              /* tp_iternext */
 760         0,                              /* tp_methods */
 761         0,                              /* tp_members */
 762         0,                              /* tp_getset */
 763         0,                              /* tp_base */
 764         0,                              /* tp_dict */
 765         0,                              /* tp_descr_get */
 766         0,                              /* tp_descr_set */
 767         0,                              /* tp_dictoffset */
 768         0,                              /* tp_init */
 769         0,                              /* tp_alloc */
 770         methodcaller_new,               /* tp_new */
 771         0,                              /* tp_free */
 772 };
 773
 774
 775 /* Initialization function for the module (*must* be called initoperator) */
 776
 777 PyMODINIT_FUNC
 778 initoperator(void)
 779 {
 780         PyObject *m;
 781
 782         /* Create the module and add the functions */
 783         m = Py_InitModule4("operator", operator_methods, operator_doc,
 784                        (PyObject*)NULL, PYTHON_API_VERSION);
 785         if (m == NULL)
 786                 return;
 787
 788         if (PyType_Ready(&itemgetter_type) < 0)
 789                 return;
 790         Py_INCREF(&itemgetter_type);
 791         PyModule_AddObject(m, "itemgetter", (PyObject *)&itemgetter_type);
 792
 793         if (PyType_Ready(&attrgetter_type) < 0)
 794                 return;
 795         Py_INCREF(&attrgetter_type);
 796         PyModule_AddObject(m, "attrgetter", (PyObject *)&attrgetter_type);
 797
 798         if (PyType_Ready(&methodcaller_type) < 0)
 799                 return;
 800         Py_INCREF(&methodcaller_type);
 801         PyModule_AddObject(m, "methodcaller", (PyObject *)&methodcaller_type);
 802 }