operator.c 31.4 KB
Newer Older
1 2 3 4 5

#include "Python.h"

PyDoc_STRVAR(operator_doc,
"Operator interface.\n\
6 7 8 9
\n\
This module exports a set of functions implemented in C corresponding\n\
to the intrinsic operators of Python.  For example, operator.add(x, y)\n\
is equivalent to the expression x+y.  The function names are those\n\
Benjamin Peterson's avatar
Benjamin Peterson committed
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 15
  return AOP(a1); }

Fred Drake's avatar
Fred Drake committed
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 20
  return AOP(a1,a2); }

Fred Drake's avatar
Fred Drake committed
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 25
  return AOP(a1,a2); }

Fred Drake's avatar
Fred Drake committed
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 30 31 32
  if(-1 == AOP(a1,a2)) return NULL; \
  Py_INCREF(Py_None); \
  return Py_None; }

Fred Drake's avatar
Fred Drake committed
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 37 38 39
  if(-1 == AOP(a1,a2,a3)) return NULL; \
  Py_INCREF(Py_None); \
  return Py_None; }

40 41
#define spami(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a1) { \
  long r; \
42
  if(-1 == (r=AOP(a1))) return NULL; \
43
  return PyBool_FromLong(r); }
44

Fred Drake's avatar
Fred Drake committed
45
#define spami2(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
46
  PyObject *a1, *a2; long r; \
47
  if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
48
  if(-1 == (r=AOP(a1,a2))) return NULL; \
49
  return PyLong_FromLong(r); }
50

51 52 53 54
#define spamn2(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
  PyObject *a1, *a2; Py_ssize_t r; \
  if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
  if(-1 == (r=AOP(a1,a2))) return NULL; \
55
  return PyLong_FromSsize_t(r); }
56

57 58
#define spami2b(OP,AOP) static PyObject *OP(PyObject *s, PyObject *a) { \
  PyObject *a1, *a2; long r; \
59
  if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
60 61 62
  if(-1 == (r=AOP(a1,a2))) return NULL; \
  return PyBool_FromLong(r); }

63 64
#define spamrc(OP,A) static PyObject *OP(PyObject *s, PyObject *a) { \
  PyObject *a1, *a2; \
65
  if(! PyArg_UnpackTuple(a,#OP,2,2,&a1,&a2)) return NULL; \
66 67
  return PyObject_RichCompare(a1,a2,A); }

68
spami(truth            , PyObject_IsTrue)
69 70 71
spam2(op_add           , PyNumber_Add)
spam2(op_sub           , PyNumber_Subtract)
spam2(op_mul           , PyNumber_Multiply)
72 73
spam2(op_floordiv      , PyNumber_FloorDivide)
spam2(op_truediv       , PyNumber_TrueDivide)
74 75 76 77 78
spam2(op_mod           , PyNumber_Remainder)
spam1(op_neg           , PyNumber_Negative)
spam1(op_pos           , PyNumber_Positive)
spam1(op_abs           , PyNumber_Absolute)
spam1(op_inv           , PyNumber_Invert)
79
spam1(op_invert        , PyNumber_Invert)
80 81
spam2(op_lshift        , PyNumber_Lshift)
spam2(op_rshift        , PyNumber_Rshift)
Guido van Rossum's avatar
Guido van Rossum committed
82
spami(op_not_          , PyObject_Not)
83 84 85
spam2(op_and_          , PyNumber_And)
spam2(op_xor           , PyNumber_Xor)
spam2(op_or_           , PyNumber_Or)
86 87 88 89 90 91 92 93 94 95 96
spam2(op_iadd          , PyNumber_InPlaceAdd)
spam2(op_isub          , PyNumber_InPlaceSubtract)
spam2(op_imul          , PyNumber_InPlaceMultiply)
spam2(op_ifloordiv     , PyNumber_InPlaceFloorDivide)
spam2(op_itruediv      , PyNumber_InPlaceTrueDivide)
spam2(op_imod          , PyNumber_InPlaceRemainder)
spam2(op_ilshift       , PyNumber_InPlaceLshift)
spam2(op_irshift       , PyNumber_InPlaceRshift)
spam2(op_iand          , PyNumber_InPlaceAnd)
spam2(op_ixor          , PyNumber_InPlaceXor)
spam2(op_ior           , PyNumber_InPlaceOr)
97
spam2(op_concat        , PySequence_Concat)
98
spam2(op_iconcat       , PySequence_InPlaceConcat)
99
spami2b(op_contains     , PySequence_Contains)
100 101
spamn2(indexOf         , PySequence_Index)
spamn2(countOf         , PySequence_Count)
102 103 104
spam2(op_getitem       , PyObject_GetItem)
spam2n(op_delitem       , PyObject_DelItem)
spam3n(op_setitem      , PyObject_SetItem)
105 106 107 108 109 110
spamrc(op_lt           , Py_LT)
spamrc(op_le           , Py_LE)
spamrc(op_eq           , Py_EQ)
spamrc(op_ne           , Py_NE)
spamrc(op_gt           , Py_GT)
spamrc(op_ge           , Py_GE)
111

112 113 114
static PyObject*
op_pow(PyObject *s, PyObject *a)
{
115 116 117 118
    PyObject *a1, *a2;
    if (PyArg_UnpackTuple(a,"pow", 2, 2, &a1, &a2))
        return PyNumber_Power(a1, a2, Py_None);
    return NULL;
119 120
}

121 122 123
static PyObject*
op_ipow(PyObject *s, PyObject *a)
{
124 125 126 127
    PyObject *a1, *a2;
    if (PyArg_UnpackTuple(a,"ipow", 2, 2, &a1, &a2))
        return PyNumber_InPlacePower(a1, a2, Py_None);
    return NULL;
128 129
}

130 131 132
static PyObject *
op_index(PyObject *s, PyObject *a)
{
133
    return PyNumber_Index(a);
134 135
}

136 137 138
static PyObject*
is_(PyObject *s, PyObject *a)
{
139 140 141 142 143 144
    PyObject *a1, *a2, *result = NULL;
    if (PyArg_UnpackTuple(a,"is_", 2, 2, &a1, &a2)) {
        result = (a1 == a2) ? Py_True : Py_False;
        Py_INCREF(result);
    }
    return result;
145 146 147 148 149
}

static PyObject*
is_not(PyObject *s, PyObject *a)
{
150 151 152 153 154 155
    PyObject *a1, *a2, *result = NULL;
    if (PyArg_UnpackTuple(a,"is_not", 2, 2, &a1, &a2)) {
        result = (a1 != a2) ? Py_True : Py_False;
        Py_INCREF(result);
    }
    return result;
156 157
}

158 159
#undef spam1
#undef spam2
160 161
#undef spam1o
#undef spam1o
162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210

/* compare_digest **********************************************************/

/*
 * timing safe compare
 *
 * Returns 1 of the strings are equal.
 * In case of len(a) != len(b) the function tries to keep the timing
 * dependent on the length of b. CPU cache locally may still alter timing
 * a bit.
 */
static int
_tscmp(const unsigned char *a, const unsigned char *b,
        Py_ssize_t len_a, Py_ssize_t len_b)
{
    /* The volatile type declarations make sure that the compiler has no
     * chance to optimize and fold the code in any way that may change
     * the timing.
     */
    volatile Py_ssize_t length;
    volatile const unsigned char *left;
    volatile const unsigned char *right;
    Py_ssize_t i;
    unsigned char result;

    /* loop count depends on length of b */
    length = len_b;
    left = NULL;
    right = b;

    /* don't use else here to keep the amount of CPU instructions constant,
     * volatile forces re-evaluation
     *  */
    if (len_a == length) {
        left = *((volatile const unsigned char**)&a);
        result = 0;
    }
    if (len_a != length) {
        left = b;
        result = 1;
    }

    for (i=0; i < length; i++) {
        result |= *left++ ^ *right++;
    }

    return (result == 0);
}

211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235
PyDoc_STRVAR(length_hint__doc__,
"length_hint(obj, default=0) -> int\n"
"Return an estimate of the number of items in obj.\n"
"This is useful for presizing containers when building from an\n"
"iterable.\n"
"\n"
"If the object supports len(), the result will be\n"
"exact. Otherwise, it may over- or under-estimate by an\n"
"arbitrary amount. The result will be an integer >= 0.");

static PyObject *length_hint(PyObject *self, PyObject *args)
{
    PyObject *obj;
    Py_ssize_t defaultvalue = 0, res;
    if (!PyArg_ParseTuple(args, "O|n:length_hint", &obj, &defaultvalue)) {
        return NULL;
    }
    res = PyObject_LengthHint(obj, defaultvalue);
    if (res == -1 && PyErr_Occurred()) {
        return NULL;
    }
    return PyLong_FromSsize_t(res);
}


236 237 238
PyDoc_STRVAR(compare_digest__doc__,
"compare_digest(a, b) -> bool\n"
"\n"
239 240 241 242
"Return 'a == b'.  This function uses an approach designed to prevent\n"
"timing analysis, making it appropriate for cryptography.\n"
"a and b must both be of the same type: either str (ASCII only),\n"
"or any type that supports the buffer protocol (e.g. bytes).\n"
243
"\n"
244
"Note: If a and b are of different lengths, or if an error occurs,\n"
245 246
"a timing attack could theoretically reveal information about the\n"
"types and lengths of a and b--but not their values.\n");
247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318

static PyObject*
compare_digest(PyObject *self, PyObject *args)
{
    PyObject *a, *b;
    int rc;

    if (!PyArg_ParseTuple(args, "OO:compare_digest", &a, &b)) {
        return NULL;
    }

    /* ASCII unicode string */
    if(PyUnicode_Check(a) && PyUnicode_Check(b)) {
        if (PyUnicode_READY(a) == -1 || PyUnicode_READY(b) == -1) {
            return NULL;
        }
        if (!PyUnicode_IS_ASCII(a) || !PyUnicode_IS_ASCII(b)) {
            PyErr_SetString(PyExc_TypeError,
                            "comparing strings with non-ASCII characters is "
                            "not supported");
            return NULL;
        }

        rc = _tscmp(PyUnicode_DATA(a),
                    PyUnicode_DATA(b),
                    PyUnicode_GET_LENGTH(a),
                    PyUnicode_GET_LENGTH(b));
    }
    /* fallback to buffer interface for bytes, bytesarray and other */
    else {
        Py_buffer view_a;
        Py_buffer view_b;

        if ((PyObject_CheckBuffer(a) == 0) & (PyObject_CheckBuffer(b) == 0)) {
            PyErr_Format(PyExc_TypeError,
                         "unsupported operand types(s) or combination of types: "
                         "'%.100s' and '%.100s'",
                         Py_TYPE(a)->tp_name, Py_TYPE(b)->tp_name);
            return NULL;
        }

        if (PyObject_GetBuffer(a, &view_a, PyBUF_SIMPLE) == -1) {
            return NULL;
        }
        if (view_a.ndim > 1) {
            PyErr_SetString(PyExc_BufferError,
                            "Buffer must be single dimension");
            PyBuffer_Release(&view_a);
            return NULL;
        }

        if (PyObject_GetBuffer(b, &view_b, PyBUF_SIMPLE) == -1) {
            PyBuffer_Release(&view_a);
            return NULL;
        }
        if (view_b.ndim > 1) {
            PyErr_SetString(PyExc_BufferError,
                            "Buffer must be single dimension");
            PyBuffer_Release(&view_a);
            PyBuffer_Release(&view_b);
            return NULL;
        }

        rc = _tscmp((const unsigned char*)view_a.buf,
                    (const unsigned char*)view_b.buf,
                    view_a.len,
                    view_b.len);

        PyBuffer_Release(&view_a);
        PyBuffer_Release(&view_b);
    }

319
    return PyBool_FromLong(rc);
320 321 322 323
}

/* operator methods **********************************************************/

324
#define spam1(OP,DOC) {#OP, OP, METH_VARARGS, PyDoc_STR(DOC)},
Armin Rigo's avatar
Armin Rigo committed
325
#define spam2(OP,ALTOP,DOC) {#OP, op_##OP, METH_VARARGS, PyDoc_STR(DOC)}, \
326
                           {#ALTOP, op_##OP, METH_VARARGS, PyDoc_STR(DOC)},
327
#define spam1o(OP,DOC) {#OP, OP, METH_O, PyDoc_STR(DOC)},
Armin Rigo's avatar
Armin Rigo committed
328
#define spam2o(OP,ALTOP,DOC) {#OP, op_##OP, METH_O, PyDoc_STR(DOC)}, \
329
                           {#ALTOP, op_##OP, METH_O, PyDoc_STR(DOC)},
330 331 332

static struct PyMethodDef operator_methods[] = {

333
spam1o(truth,
334
 "truth(a) -- Return True if a is true, False otherwise.")
335
spam2(contains,__contains__,
336
 "contains(a, b) -- Same as b in a (note reversed operands).")
337
spam1(indexOf,
338
 "indexOf(a, b) -- Return the first index of b in a.")
339 340
spam1(countOf,
 "countOf(a, b) -- Return the number of times b occurs in a.")
341

342 343
spam1(is_, "is_(a, b) -- Same as a is b.")
spam1(is_not, "is_not(a, b) -- Same as a is not b.")
344
spam2o(index, __index__, "index(a) -- Same as a.__index__()")
345 346 347
spam2(add,__add__, "add(a, b) -- Same as a + b.")
spam2(sub,__sub__, "sub(a, b) -- Same as a - b.")
spam2(mul,__mul__, "mul(a, b) -- Same as a * b.")
348
spam2(floordiv,__floordiv__, "floordiv(a, b) -- Same as a // b.")
349
spam2(truediv,__truediv__, "truediv(a, b) -- Same as a / b.")
350
spam2(mod,__mod__, "mod(a, b) -- Same as a % b.")
351 352 353 354 355
spam2o(neg,__neg__, "neg(a) -- Same as -a.")
spam2o(pos,__pos__, "pos(a) -- Same as +a.")
spam2o(abs,__abs__, "abs(a) -- Same as abs(a).")
spam2o(inv,__inv__, "inv(a) -- Same as ~a.")
spam2o(invert,__invert__, "invert(a) -- Same as ~a.")
356 357
spam2(lshift,__lshift__, "lshift(a, b) -- Same as a << b.")
spam2(rshift,__rshift__, "rshift(a, b) -- Same as a >> b.")
358
spam2o(not_,__not__, "not_(a) -- Same as not a.")
359 360 361
spam2(and_,__and__, "and_(a, b) -- Same as a & b.")
spam2(xor,__xor__, "xor(a, b) -- Same as a ^ b.")
spam2(or_,__or__, "or_(a, b) -- Same as a | b.")
Benjamin Peterson's avatar
Benjamin Peterson committed
362 363 364 365 366 367 368 369 370 371 372
spam2(iadd,__iadd__, "a = iadd(a, b) -- Same as a += b.")
spam2(isub,__isub__, "a = isub(a, b) -- Same as a -= b.")
spam2(imul,__imul__, "a = imul(a, b) -- Same as a *= b.")
spam2(ifloordiv,__ifloordiv__, "a = ifloordiv(a, b) -- Same as a //= b.")
spam2(itruediv,__itruediv__, "a = itruediv(a, b) -- Same as a /= b")
spam2(imod,__imod__, "a = imod(a, b) -- Same as a %= b.")
spam2(ilshift,__ilshift__, "a = ilshift(a, b) -- Same as a <<= b.")
spam2(irshift,__irshift__, "a = irshift(a, b) -- Same as a >>= b.")
spam2(iand,__iand__, "a = iand(a, b) -- Same as a &= b.")
spam2(ixor,__ixor__, "a = ixor(a, b) -- Same as a ^= b.")
spam2(ior,__ior__, "a = ior(a, b) -- Same as a |= b.")
373
spam2(concat,__concat__,
374
 "concat(a, b) -- Same as a + b, for a and b sequences.")
375
spam2(iconcat,__iconcat__,
Benjamin Peterson's avatar
Benjamin Peterson committed
376
 "a = iconcat(a, b) -- Same as a += b, for a and b sequences.")
377
spam2(getitem,__getitem__,
378
 "getitem(a, b) -- Same as a[b].")
379
spam2(setitem,__setitem__,
380
 "setitem(a, b, c) -- Same as a[b] = c.")
381
spam2(delitem,__delitem__,
382
 "delitem(a, b) -- Same as del a[b].")
383
spam2(pow,__pow__, "pow(a, b) -- Same as a ** b.")
Benjamin Peterson's avatar
Benjamin Peterson committed
384
spam2(ipow,__ipow__, "a = ipow(a, b) -- Same as a **= b.")
385 386 387 388 389 390
spam2(lt,__lt__, "lt(a, b) -- Same as a<b.")
spam2(le,__le__, "le(a, b) -- Same as a<=b.")
spam2(eq,__eq__, "eq(a, b) -- Same as a==b.")
spam2(ne,__ne__, "ne(a, b) -- Same as a!=b.")
spam2(gt,__gt__, "gt(a, b) -- Same as a>b.")
spam2(ge,__ge__, "ge(a, b) -- Same as a>=b.")
391

392 393
    {"_compare_digest", (PyCFunction)compare_digest, METH_VARARGS,
     compare_digest__doc__},
394 395
     {"length_hint", (PyCFunction)length_hint, METH_VARARGS,
     length_hint__doc__},
396
    {NULL,              NULL}           /* sentinel */
397 398 399

};

400
/* itemgetter object **********************************************************/
401

402
typedef struct {
403 404 405
    PyObject_HEAD
    Py_ssize_t nitems;
    PyObject *item;
406 407 408 409 410 411 412
} itemgetterobject;

static PyTypeObject itemgetter_type;

static PyObject *
itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437
    itemgetterobject *ig;
    PyObject *item;
    Py_ssize_t nitems;

    if (!_PyArg_NoKeywords("itemgetter()", kwds))
        return NULL;

    nitems = PyTuple_GET_SIZE(args);
    if (nitems <= 1) {
        if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &item))
            return NULL;
    } else
        item = args;

    /* create itemgetterobject structure */
    ig = PyObject_GC_New(itemgetterobject, &itemgetter_type);
    if (ig == NULL)
        return NULL;

    Py_INCREF(item);
    ig->item = item;
    ig->nitems = nitems;

    PyObject_GC_Track(ig);
    return (PyObject *)ig;
438 439 440 441 442
}

static void
itemgetter_dealloc(itemgetterobject *ig)
{
443 444 445
    PyObject_GC_UnTrack(ig);
    Py_XDECREF(ig->item);
    PyObject_GC_Del(ig);
446 447 448 449 450
}

static int
itemgetter_traverse(itemgetterobject *ig, visitproc visit, void *arg)
{
451 452
    Py_VISIT(ig->item);
    return 0;
453 454 455 456 457
}

static PyObject *
itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw)
{
458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483
    PyObject *obj, *result;
    Py_ssize_t i, nitems=ig->nitems;

    if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj))
        return NULL;
    if (nitems == 1)
        return PyObject_GetItem(obj, ig->item);

    assert(PyTuple_Check(ig->item));
    assert(PyTuple_GET_SIZE(ig->item) == nitems);

    result = PyTuple_New(nitems);
    if (result == NULL)
        return NULL;

    for (i=0 ; i < nitems ; i++) {
        PyObject *item, *val;
        item = PyTuple_GET_ITEM(ig->item, i);
        val = PyObject_GetItem(obj, item);
        if (val == NULL) {
            Py_DECREF(result);
            return NULL;
        }
        PyTuple_SET_ITEM(result, i, val);
    }
    return result;
484 485 486
}

PyDoc_STRVAR(itemgetter_doc,
487
"itemgetter(item, ...) --> itemgetter object\n\
488
\n\
489 490 491
Return a callable object that fetches the given item(s) from its operand.\n\
After, f=itemgetter(2), the call f(r) returns r[2].\n\
After, g=itemgetter(2,5,3), the call g(r) returns (r[2], r[5], r[3])");
492 493

static PyTypeObject itemgetter_type = {
494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533
    PyVarObject_HEAD_INIT(NULL, 0)
    "operator.itemgetter",              /* tp_name */
    sizeof(itemgetterobject),           /* tp_basicsize */
    0,                                  /* tp_itemsize */
    /* methods */
    (destructor)itemgetter_dealloc,     /* tp_dealloc */
    0,                                  /* tp_print */
    0,                                  /* tp_getattr */
    0,                                  /* tp_setattr */
    0,                                  /* tp_reserved */
    0,                                  /* tp_repr */
    0,                                  /* tp_as_number */
    0,                                  /* tp_as_sequence */
    0,                                  /* tp_as_mapping */
    0,                                  /* tp_hash */
    (ternaryfunc)itemgetter_call,       /* tp_call */
    0,                                  /* tp_str */
    PyObject_GenericGetAttr,            /* tp_getattro */
    0,                                  /* tp_setattro */
    0,                                  /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,            /* tp_flags */
    itemgetter_doc,                     /* tp_doc */
    (traverseproc)itemgetter_traverse,          /* tp_traverse */
    0,                                  /* tp_clear */
    0,                                  /* tp_richcompare */
    0,                                  /* tp_weaklistoffset */
    0,                                  /* tp_iter */
    0,                                  /* tp_iternext */
    0,                                  /* tp_methods */
    0,                                  /* tp_members */
    0,                                  /* tp_getset */
    0,                                  /* tp_base */
    0,                                  /* tp_dict */
    0,                                  /* tp_descr_get */
    0,                                  /* tp_descr_set */
    0,                                  /* tp_dictoffset */
    0,                                  /* tp_init */
    0,                                  /* tp_alloc */
    itemgetter_new,                     /* tp_new */
    0,                                  /* tp_free */
534 535 536 537 538 539
};


/* attrgetter object **********************************************************/

typedef struct {
540 541 542
    PyObject_HEAD
    Py_ssize_t nattrs;
    PyObject *attr;
543 544 545 546 547 548 549
} attrgetterobject;

static PyTypeObject attrgetter_type;

static PyObject *
attrgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
550 551
    attrgetterobject *ag;
    PyObject *attr;
552
    Py_ssize_t nattrs, idx, char_idx;
553 554 555 556 557 558 559 560

    if (!_PyArg_NoKeywords("attrgetter()", kwds))
        return NULL;

    nattrs = PyTuple_GET_SIZE(args);
    if (nattrs <= 1) {
        if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &attr))
            return NULL;
561 562 563 564 565 566 567 568 569 570
    }

    attr = PyTuple_New(nattrs);
    if (attr == NULL)
        return NULL;

    /* prepare attr while checking args */
    for (idx = 0; idx < nattrs; ++idx) {
        PyObject *item = PyTuple_GET_ITEM(args, idx);
        Py_ssize_t item_len;
Martin v. Löwis's avatar
Martin v. Löwis committed
571 572
        void *data;
        unsigned int kind;
573 574 575 576 577 578 579 580
        int dot_count;

        if (!PyUnicode_Check(item)) {
            PyErr_SetString(PyExc_TypeError,
                            "attribute name must be a string");
            Py_DECREF(attr);
            return NULL;
        }
Martin v. Löwis's avatar
Martin v. Löwis committed
581 582 583 584 585 586 587
        if (PyUnicode_READY(item)) {
            Py_DECREF(attr);
            return NULL;
        }
        item_len = PyUnicode_GET_LENGTH(item);
        kind = PyUnicode_KIND(item);
        data = PyUnicode_DATA(item);
588 589 590 591

        /* check whethere the string is dotted */
        dot_count = 0;
        for (char_idx = 0; char_idx < item_len; ++char_idx) {
Martin v. Löwis's avatar
Martin v. Löwis committed
592
            if (PyUnicode_READ(kind, data, char_idx) == '.')
593 594 595 596 597 598 599 600 601 602
                ++dot_count;
        }

        if (dot_count == 0) {
            Py_INCREF(item);
            PyUnicode_InternInPlace(&item);
            PyTuple_SET_ITEM(attr, idx, item);
        } else { /* make it a tuple of non-dotted attrnames */
            PyObject *attr_chain = PyTuple_New(dot_count + 1);
            PyObject *attr_chain_item;
603 604 605
            Py_ssize_t unibuff_from = 0;
            Py_ssize_t unibuff_till = 0;
            Py_ssize_t attr_chain_idx = 0;
606 607 608 609 610 611 612

            if (attr_chain == NULL) {
                Py_DECREF(attr);
                return NULL;
            }

            for (; dot_count > 0; --dot_count) {
Martin v. Löwis's avatar
Martin v. Löwis committed
613
                while (PyUnicode_READ(kind, data, unibuff_till) != '.') {
614 615
                    ++unibuff_till;
                }
Martin v. Löwis's avatar
Martin v. Löwis committed
616 617 618
                attr_chain_item = PyUnicode_Substring(item,
                                      unibuff_from,
                                      unibuff_till);
619 620 621 622 623 624 625 626 627 628 629 630
                if (attr_chain_item == NULL) {
                    Py_DECREF(attr_chain);
                    Py_DECREF(attr);
                    return NULL;
                }
                PyUnicode_InternInPlace(&attr_chain_item);
                PyTuple_SET_ITEM(attr_chain, attr_chain_idx, attr_chain_item);
                ++attr_chain_idx;
                unibuff_till = unibuff_from = unibuff_till + 1;
            }

            /* now add the last dotless name */
Martin v. Löwis's avatar
Martin v. Löwis committed
631 632
            attr_chain_item = PyUnicode_Substring(item,
                                                  unibuff_from, item_len);
633 634 635 636 637 638 639 640 641 642 643
            if (attr_chain_item == NULL) {
                Py_DECREF(attr_chain);
                Py_DECREF(attr);
                return NULL;
            }
            PyUnicode_InternInPlace(&attr_chain_item);
            PyTuple_SET_ITEM(attr_chain, attr_chain_idx, attr_chain_item);

            PyTuple_SET_ITEM(attr, idx, attr_chain);
        }
    }
644 645 646

    /* create attrgetterobject structure */
    ag = PyObject_GC_New(attrgetterobject, &attrgetter_type);
647 648
    if (ag == NULL) {
        Py_DECREF(attr);
649
        return NULL;
650
    }
651 652 653 654 655 656

    ag->attr = attr;
    ag->nattrs = nattrs;

    PyObject_GC_Track(ag);
    return (PyObject *)ag;
657 658 659 660 661
}

static void
attrgetter_dealloc(attrgetterobject *ag)
{
662 663 664
    PyObject_GC_UnTrack(ag);
    Py_XDECREF(ag->attr);
    PyObject_GC_Del(ag);
665 666 667 668 669
}

static int
attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg)
{
670 671
    Py_VISIT(ag->attr);
    return 0;
672 673
}

674 675 676
static PyObject *
dotted_getattr(PyObject *obj, PyObject *attr)
{
677 678 679 680 681 682 683 684 685 686 687 688 689
    PyObject *newobj;

    /* attr is either a tuple or instance of str.
       Ensured by the setup code of attrgetter_new */
    if (PyTuple_CheckExact(attr)) { /* chained getattr */
        Py_ssize_t name_idx = 0, name_count;
        PyObject *attr_name;

        name_count = PyTuple_GET_SIZE(attr);
        Py_INCREF(obj);
        for (name_idx = 0; name_idx < name_count; ++name_idx) {
            attr_name = PyTuple_GET_ITEM(attr, name_idx);
            newobj = PyObject_GetAttr(obj, attr_name);
690
            Py_DECREF(obj);
691 692 693 694 695
            if (newobj == NULL) {
                return NULL;
            }
            /* here */
            obj = newobj;
696
        }
697 698
    } else { /* single getattr */
        newobj = PyObject_GetAttr(obj, attr);
699 700 701 702 703 704
        if (newobj == NULL)
            return NULL;
        obj = newobj;
    }

    return obj;
705 706
}

707 708 709
static PyObject *
attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
{
710 711 712 713 714
    PyObject *obj, *result;
    Py_ssize_t i, nattrs=ag->nattrs;

    if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj))
        return NULL;
715 716
    if (ag->nattrs == 1) /* ag->attr is always a tuple */
        return dotted_getattr(obj, PyTuple_GET_ITEM(ag->attr, 0));
717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735

    assert(PyTuple_Check(ag->attr));
    assert(PyTuple_GET_SIZE(ag->attr) == nattrs);

    result = PyTuple_New(nattrs);
    if (result == NULL)
        return NULL;

    for (i=0 ; i < nattrs ; i++) {
        PyObject *attr, *val;
        attr = PyTuple_GET_ITEM(ag->attr, i);
        val = dotted_getattr(obj, attr);
        if (val == NULL) {
            Py_DECREF(result);
            return NULL;
        }
        PyTuple_SET_ITEM(result, i, val);
    }
    return result;
736 737 738
}

PyDoc_STRVAR(attrgetter_doc,
739
"attrgetter(attr, ...) --> attrgetter object\n\
740
\n\
741 742
Return a callable object that fetches the given attribute(s) from its operand.\n\
After, f=attrgetter('name'), the call f(r) returns r.name.\n\
743 744 745
After, g=attrgetter('name', 'date'), the call g(r) returns (r.name, r.date).\n\
After, h=attrgetter('name.first', 'name.last'), the call h(r) returns\n\
(r.name.first, r.name.last).");
746 747

static PyTypeObject attrgetter_type = {
748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787
    PyVarObject_HEAD_INIT(NULL, 0)
    "operator.attrgetter",              /* tp_name */
    sizeof(attrgetterobject),           /* tp_basicsize */
    0,                                  /* tp_itemsize */
    /* methods */
    (destructor)attrgetter_dealloc,     /* tp_dealloc */
    0,                                  /* tp_print */
    0,                                  /* tp_getattr */
    0,                                  /* tp_setattr */
    0,                                  /* tp_reserved */
    0,                                  /* tp_repr */
    0,                                  /* tp_as_number */
    0,                                  /* tp_as_sequence */
    0,                                  /* tp_as_mapping */
    0,                                  /* tp_hash */
    (ternaryfunc)attrgetter_call,       /* tp_call */
    0,                                  /* tp_str */
    PyObject_GenericGetAttr,            /* tp_getattro */
    0,                                  /* tp_setattro */
    0,                                  /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,            /* tp_flags */
    attrgetter_doc,                     /* tp_doc */
    (traverseproc)attrgetter_traverse,          /* tp_traverse */
    0,                                  /* tp_clear */
    0,                                  /* tp_richcompare */
    0,                                  /* tp_weaklistoffset */
    0,                                  /* tp_iter */
    0,                                  /* tp_iternext */
    0,                                  /* tp_methods */
    0,                                  /* tp_members */
    0,                                  /* tp_getset */
    0,                                  /* tp_base */
    0,                                  /* tp_dict */
    0,                                  /* tp_descr_get */
    0,                                  /* tp_descr_set */
    0,                                  /* tp_dictoffset */
    0,                                  /* tp_init */
    0,                                  /* tp_alloc */
    attrgetter_new,                     /* tp_new */
    0,                                  /* tp_free */
788
};
789 790 791 792 793


/* methodcaller object **********************************************************/

typedef struct {
794 795 796 797
    PyObject_HEAD
    PyObject *name;
    PyObject *args;
    PyObject *kwds;
798 799 800 801 802 803 804
} methodcallerobject;

static PyTypeObject methodcaller_type;

static PyObject *
methodcaller_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834
    methodcallerobject *mc;
    PyObject *name, *newargs;

    if (PyTuple_GET_SIZE(args) < 1) {
        PyErr_SetString(PyExc_TypeError, "methodcaller needs at least "
                        "one argument, the method name");
        return NULL;
    }

    /* create methodcallerobject structure */
    mc = PyObject_GC_New(methodcallerobject, &methodcaller_type);
    if (mc == NULL)
        return NULL;

    newargs = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
    if (newargs == NULL) {
        Py_DECREF(mc);
        return NULL;
    }
    mc->args = newargs;

    name = PyTuple_GET_ITEM(args, 0);
    Py_INCREF(name);
    mc->name = name;

    Py_XINCREF(kwds);
    mc->kwds = kwds;

    PyObject_GC_Track(mc);
    return (PyObject *)mc;
835 836 837 838 839
}

static void
methodcaller_dealloc(methodcallerobject *mc)
{
840 841 842 843 844
    PyObject_GC_UnTrack(mc);
    Py_XDECREF(mc->name);
    Py_XDECREF(mc->args);
    Py_XDECREF(mc->kwds);
    PyObject_GC_Del(mc);
845 846 847 848 849
}

static int
methodcaller_traverse(methodcallerobject *mc, visitproc visit, void *arg)
{
850 851 852
    Py_VISIT(mc->args);
    Py_VISIT(mc->kwds);
    return 0;
853 854 855 856 857
}

static PyObject *
methodcaller_call(methodcallerobject *mc, PyObject *args, PyObject *kw)
{
858 859 860 861 862 863 864 865 866 867
    PyObject *method, *obj, *result;

    if (!PyArg_UnpackTuple(args, "methodcaller", 1, 1, &obj))
        return NULL;
    method = PyObject_GetAttr(obj, mc->name);
    if (method == NULL)
        return NULL;
    result = PyObject_Call(method, mc->args, mc->kwds);
    Py_DECREF(method);
    return result;
868 869 870 871 872 873 874 875 876 877 878
}

PyDoc_STRVAR(methodcaller_doc,
"methodcaller(name, ...) --> methodcaller object\n\
\n\
Return a callable object that calls the given method on its operand.\n\
After, f = methodcaller('name'), the call f(r) returns r.name().\n\
After, g = methodcaller('name', 'date', foo=1), the call g(r) returns\n\
r.name('date', foo=1).");

static PyTypeObject methodcaller_type = {
879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918
    PyVarObject_HEAD_INIT(NULL, 0)
    "operator.methodcaller",            /* tp_name */
    sizeof(methodcallerobject),         /* tp_basicsize */
    0,                                  /* tp_itemsize */
    /* methods */
    (destructor)methodcaller_dealloc, /* tp_dealloc */
    0,                                  /* tp_print */
    0,                                  /* tp_getattr */
    0,                                  /* tp_setattr */
    0,                                  /* tp_reserved */
    0,                                  /* tp_repr */
    0,                                  /* tp_as_number */
    0,                                  /* tp_as_sequence */
    0,                                  /* tp_as_mapping */
    0,                                  /* tp_hash */
    (ternaryfunc)methodcaller_call,     /* tp_call */
    0,                                  /* tp_str */
    PyObject_GenericGetAttr,            /* tp_getattro */
    0,                                  /* tp_setattro */
    0,                                  /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
    methodcaller_doc,                           /* tp_doc */
    (traverseproc)methodcaller_traverse,        /* tp_traverse */
    0,                                  /* tp_clear */
    0,                                  /* tp_richcompare */
    0,                                  /* tp_weaklistoffset */
    0,                                  /* tp_iter */
    0,                                  /* tp_iternext */
    0,                                  /* tp_methods */
    0,                                  /* tp_members */
    0,                                  /* tp_getset */
    0,                                  /* tp_base */
    0,                                  /* tp_dict */
    0,                                  /* tp_descr_get */
    0,                                  /* tp_descr_set */
    0,                                  /* tp_dictoffset */
    0,                                  /* tp_init */
    0,                                  /* tp_alloc */
    methodcaller_new,                   /* tp_new */
    0,                                  /* tp_free */
919 920 921
};


922 923 924 925
/* Initialization function for the module (*must* be called PyInit_operator) */


static struct PyModuleDef operatormodule = {
926 927 928 929 930 931 932 933 934
    PyModuleDef_HEAD_INIT,
    "operator",
    operator_doc,
    -1,
    operator_methods,
    NULL,
    NULL,
    NULL,
    NULL
935
};
936

937
PyMODINIT_FUNC
938
PyInit_operator(void)
939
{
940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961
    PyObject *m;

    /* Create the module and add the functions */
    m = PyModule_Create(&operatormodule);
    if (m == NULL)
        return NULL;

    if (PyType_Ready(&itemgetter_type) < 0)
        return NULL;
    Py_INCREF(&itemgetter_type);
    PyModule_AddObject(m, "itemgetter", (PyObject *)&itemgetter_type);

    if (PyType_Ready(&attrgetter_type) < 0)
        return NULL;
    Py_INCREF(&attrgetter_type);
    PyModule_AddObject(m, "attrgetter", (PyObject *)&attrgetter_type);

    if (PyType_Ready(&methodcaller_type) < 0)
        return NULL;
    Py_INCREF(&methodcaller_type);
    PyModule_AddObject(m, "methodcaller", (PyObject *)&methodcaller_type);
    return m;
962
}