Modules/_functoolsmodule.c

   1
   2 #include "Python.h"
   3 #include "structmember.h"
   4
   5 /* _functools module written and maintained
   6    by Hye-Shik Chang <perky@FreeBSD.org>
   7    with adaptations by Raymond Hettinger <python@rcn.com>
   8    Copyright (c) 2004, 2005, 2006 Python Software Foundation.
   9    All rights reserved.
  10 */
  11
  12 /* reduce() *************************************************************/
  13
  14 static PyObject *
  15 functools_reduce(PyObject *self, PyObject *args)
  16 {
  17         PyObject *seq, *func, *result = NULL, *it;
  18
  19         if (!PyArg_UnpackTuple(args, "reduce", 2, 3, &func, &seq, &result))
  20                 return NULL;
  21         if (result != NULL)
  22                 Py_INCREF(result);
  23
  24         it = PyObject_GetIter(seq);
  25         if (it == NULL) {
  26                 PyErr_SetString(PyExc_TypeError,
  27                     "reduce() arg 2 must support iteration");
  28                 Py_XDECREF(result);
  29                 return NULL;
  30         }
  31
  32         if ((args = PyTuple_New(2)) == NULL)
  33                 goto Fail;
  34
  35         for (;;) {
  36                 PyObject *op2;
  37
  38                 if (args->ob_refcnt > 1) {
  39                         Py_DECREF(args);
  40                         if ((args = PyTuple_New(2)) == NULL)
  41                                 goto Fail;
  42                 }
  43
  44                 op2 = PyIter_Next(it);
  45                 if (op2 == NULL) {
  46                         if (PyErr_Occurred())
  47                                 goto Fail;
  48                         break;
  49                 }
  50
  51                 if (result == NULL)
  52                         result = op2;
  53                 else {
  54                         PyTuple_SetItem(args, 0, result);
  55                         PyTuple_SetItem(args, 1, op2);
  56                         if ((result = PyEval_CallObject(func, args)) == NULL)
  57                                 goto Fail;
  58                 }
  59         }
  60
  61         Py_DECREF(args);
  62
  63         if (result == NULL)
  64                 PyErr_SetString(PyExc_TypeError,
  65                            "reduce() of empty sequence with no initial value");
  66
  67         Py_DECREF(it);
  68         return result;
  69
  70 Fail:
  71         Py_XDECREF(args);
  72         Py_XDECREF(result);
  73         Py_DECREF(it);
  74         return NULL;
  75 }
  76
  77 PyDoc_STRVAR(reduce_doc,
  78 "reduce(function, sequence[, initial]) -> value\n\
  79 \n\
  80 Apply a function of two arguments cumulatively to the items of a sequence,\n\
  81 from left to right, so as to reduce the sequence to a single value.\n\
  82 For example, reduce(lambda x, y: x+y, [1, 2, 3, 4, 5]) calculates\n\
  83 ((((1+2)+3)+4)+5).  If initial is present, it is placed before the items\n\
  84 of the sequence in the calculation, and serves as a default when the\n\
  85 sequence is empty.");
  86
  87
  88
  89
  90 /* partial object **********************************************************/
  91
  92 typedef struct {
  93         PyObject_HEAD
  94         PyObject *fn;
  95         PyObject *args;
  96         PyObject *kw;
  97         PyObject *dict;
  98         PyObject *weakreflist; /* List of weak references */
  99 } partialobject;
 100
 101 static PyTypeObject partial_type;
 102
 103 static PyObject *
 104 partial_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 105 {
 106         PyObject *func;
 107         partialobject *pto;
 108
 109         if (PyTuple_GET_SIZE(args) < 1) {
 110                 PyErr_SetString(PyExc_TypeError,
 111                                 "type 'partial' takes at least one argument");
 112                 return NULL;
 113         }
 114
 115         func = PyTuple_GET_ITEM(args, 0);
 116         if (!PyCallable_Check(func)) {
 117                 PyErr_SetString(PyExc_TypeError,
 118                                 "the first argument must be callable");
 119                 return NULL;
 120         }
 121
 122         /* create partialobject structure */
 123         pto = (partialobject *)type->tp_alloc(type, 0);
 124         if (pto == NULL)
 125                 return NULL;
 126
 127         pto->fn = func;
 128         Py_INCREF(func);
 129         pto->args = PyTuple_GetSlice(args, 1, PY_SSIZE_T_MAX);
 130         if (pto->args == NULL) {
 131                 pto->kw = NULL;
 132                 Py_DECREF(pto);
 133                 return NULL;
 134         }
 135         if (kw != NULL) {
 136                 pto->kw = PyDict_Copy(kw);
 137                 if (pto->kw == NULL) {
 138                         Py_DECREF(pto);
 139                         return NULL;
 140                 }
 141         } else {
 142                 pto->kw = Py_None;
 143                 Py_INCREF(Py_None);
 144         }
 145
 146         pto->weakreflist = NULL;
 147         pto->dict = NULL;
 148
 149         return (PyObject *)pto;
 150 }
 151
 152 static void
 153 partial_dealloc(partialobject *pto)
 154 {
 155         PyObject_GC_UnTrack(pto);
 156         if (pto->weakreflist != NULL)
 157                 PyObject_ClearWeakRefs((PyObject *) pto);
 158         Py_XDECREF(pto->fn);
 159         Py_XDECREF(pto->args);
 160         Py_XDECREF(pto->kw);
 161         Py_XDECREF(pto->dict);
 162         Py_TYPE(pto)->tp_free(pto);
 163 }
 164
 165 static PyObject *
 166 partial_call(partialobject *pto, PyObject *args, PyObject *kw)
 167 {
 168         PyObject *ret;
 169         PyObject *argappl = NULL, *kwappl = NULL;
 170
 171         assert (PyCallable_Check(pto->fn));
 172         assert (PyTuple_Check(pto->args));
 173         assert (pto->kw == Py_None  ||  PyDict_Check(pto->kw));
 174
 175         if (PyTuple_GET_SIZE(pto->args) == 0) {
 176                 argappl = args;
 177                 Py_INCREF(args);
 178         } else if (PyTuple_GET_SIZE(args) == 0) {
 179                 argappl = pto->args;
 180                 Py_INCREF(pto->args);
 181         } else {
 182                 argappl = PySequence_Concat(pto->args, args);
 183                 if (argappl == NULL)
 184                         return NULL;
 185         }
 186
 187         if (pto->kw == Py_None) {
 188                 kwappl = kw;
 189                 Py_XINCREF(kw);
 190         } else {
 191                 kwappl = PyDict_Copy(pto->kw);
 192                 if (kwappl == NULL) {
 193                         Py_DECREF(argappl);
 194                         return NULL;
 195                 }
 196                 if (kw != NULL) {
 197                         if (PyDict_Merge(kwappl, kw, 1) != 0) {
 198                                 Py_DECREF(argappl);
 199                                 Py_DECREF(kwappl);
 200                                 return NULL;
 201                         }
 202                 }
 203         }
 204
 205         ret = PyObject_Call(pto->fn, argappl, kwappl);
 206         Py_DECREF(argappl);
 207         Py_XDECREF(kwappl);
 208         return ret;
 209 }
 210
 211 static int
 212 partial_traverse(partialobject *pto, visitproc visit, void *arg)
 213 {
 214         Py_VISIT(pto->fn);
 215         Py_VISIT(pto->args);
 216         Py_VISIT(pto->kw);
 217         Py_VISIT(pto->dict);
 218         return 0;
 219 }
 220
 221 PyDoc_STRVAR(partial_doc,
 222 "partial(func, *args, **keywords) - new function with partial application\n\
 223         of the given arguments and keywords.\n");
 224
 225 #define OFF(x) offsetof(partialobject, x)
 226 static PyMemberDef partial_memberlist[] = {
 227         {"func",        T_OBJECT,       OFF(fn),        READONLY,
 228          "function object to use in future partial calls"},
 229         {"args",        T_OBJECT,       OFF(args),      READONLY,
 230          "tuple of arguments to future partial calls"},
 231         {"keywords",    T_OBJECT,       OFF(kw),        READONLY,
 232          "dictionary of keyword arguments to future partial calls"},
 233         {NULL}  /* Sentinel */
 234 };
 235
 236 static PyObject *
 237 partial_get_dict(partialobject *pto)
 238 {
 239         if (pto->dict == NULL) {
 240                 pto->dict = PyDict_New();
 241                 if (pto->dict == NULL)
 242                         return NULL;
 243         }
 244         Py_INCREF(pto->dict);
 245         return pto->dict;
 246 }
 247
 248 static int
 249 partial_set_dict(partialobject *pto, PyObject *value)
 250 {
 251         PyObject *tmp;
 252
 253         /* It is illegal to del p.__dict__ */
 254         if (value == NULL) {
 255                 PyErr_SetString(PyExc_TypeError,
 256                                 "a partial object's dictionary may not be deleted");
 257                 return -1;
 258         }
 259         /* Can only set __dict__ to a dictionary */
 260         if (!PyDict_Check(value)) {
 261                 PyErr_SetString(PyExc_TypeError,
 262                                 "setting partial object's dictionary to a non-dict");
 263                 return -1;
 264         }
 265         tmp = pto->dict;
 266         Py_INCREF(value);
 267         pto->dict = value;
 268         Py_XDECREF(tmp);
 269         return 0;
 270 }
 271
 272 static PyGetSetDef partial_getsetlist[] = {
 273         {"__dict__", (getter)partial_get_dict, (setter)partial_set_dict},
 274         {NULL} /* Sentinel */
 275 };
 276
 277 /* Pickle strategy:
 278    __reduce__ by itself doesn't support getting kwargs in the unpickle
 279    operation so we define a __setstate__ that replaces all the information
 280    about the partial.  If we only replaced part of it someone would use
 281    it as a hook to do stange things.
 282  */
 283
 284 PyObject *
 285 partial_reduce(partialobject *pto, PyObject *unused)
 286 {
 287         return Py_BuildValue("O(O)(OOOO)", Py_TYPE(pto), pto->fn, pto->fn,
 288                              pto->args, pto->kw,
 289                              pto->dict ? pto->dict : Py_None);
 290 }
 291
 292 PyObject *
 293 partial_setstate(partialobject *pto, PyObject *args)
 294 {
 295         PyObject *fn, *fnargs, *kw, *dict;
 296         if (!PyArg_ParseTuple(args, "(OOOO):__setstate__",
 297                               &fn, &fnargs, &kw, &dict))
 298                 return NULL;
 299         Py_XDECREF(pto->fn);
 300         Py_XDECREF(pto->args);
 301         Py_XDECREF(pto->kw);
 302         Py_XDECREF(pto->dict);
 303         pto->fn = fn;
 304         pto->args = fnargs;
 305         pto->kw = kw;
 306         if (dict != Py_None) {
 307           pto->dict = dict;
 308           Py_INCREF(dict);
 309         } else {
 310           pto->dict = NULL;
 311         }
 312         Py_INCREF(fn);
 313         Py_INCREF(fnargs);
 314         Py_INCREF(kw);
 315         Py_RETURN_NONE;
 316 }
 317
 318 static PyMethodDef partial_methods[] = {
 319         {"__reduce__", (PyCFunction)partial_reduce, METH_NOARGS},
 320         {"__setstate__", (PyCFunction)partial_setstate, METH_VARARGS},
 321         {NULL,          NULL}           /* sentinel */
 322 };
 323
 324 static PyTypeObject partial_type = {
 325         PyVarObject_HEAD_INIT(NULL, 0)
 326         "functools.partial",            /* tp_name */
 327         sizeof(partialobject),          /* tp_basicsize */
 328         0,                              /* tp_itemsize */
 329         /* methods */
 330         (destructor)partial_dealloc,    /* tp_dealloc */
 331         0,                              /* tp_print */
 332         0,                              /* tp_getattr */
 333         0,                              /* tp_setattr */
 334         0,                              /* tp_compare */
 335         0,                              /* tp_repr */
 336         0,                              /* tp_as_number */
 337         0,                              /* tp_as_sequence */
 338         0,                              /* tp_as_mapping */
 339         0,                              /* tp_hash */
 340         (ternaryfunc)partial_call,      /* tp_call */
 341         0,                              /* tp_str */
 342         PyObject_GenericGetAttr,        /* tp_getattro */
 343         PyObject_GenericSetAttr,        /* tp_setattro */
 344         0,                              /* tp_as_buffer */
 345         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 346                 Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_WEAKREFS, /* tp_flags */
 347         partial_doc,                    /* tp_doc */
 348         (traverseproc)partial_traverse, /* tp_traverse */
 349         0,                              /* tp_clear */
 350         0,                              /* tp_richcompare */
 351         offsetof(partialobject, weakreflist),   /* tp_weaklistoffset */
 352         0,                              /* tp_iter */
 353         0,                              /* tp_iternext */
 354         partial_methods,                /* tp_methods */
 355         partial_memberlist,             /* tp_members */
 356         partial_getsetlist,             /* tp_getset */
 357         0,                              /* tp_base */
 358         0,                              /* tp_dict */
 359         0,                              /* tp_descr_get */
 360         0,                              /* tp_descr_set */
 361         offsetof(partialobject, dict),  /* tp_dictoffset */
 362         0,                              /* tp_init */
 363         0,                              /* tp_alloc */
 364         partial_new,                    /* tp_new */
 365         PyObject_GC_Del,                /* tp_free */
 366 };
 367
 368
 369 /* module level code ********************************************************/
 370
 371 PyDoc_STRVAR(module_doc,
 372 "Tools that operate on functions.");
 373
 374 static PyMethodDef module_methods[] = {
 375         {"reduce",      functools_reduce,     METH_VARARGS, reduce_doc},
 376         {NULL,          NULL}           /* sentinel */
 377 };
 378
 379 PyMODINIT_FUNC
 380 init_functools(void)
 381 {
 382         int i;
 383         PyObject *m;
 384         char *name;
 385         PyTypeObject *typelist[] = {
 386                 &partial_type,
 387                 NULL
 388         };
 389
 390         m = Py_InitModule3("_functools", module_methods, module_doc);
 391         if (m == NULL)
 392                 return;
 393
 394         for (i=0 ; typelist[i] != NULL ; i++) {
 395                 if (PyType_Ready(typelist[i]) < 0)
 396                         return;
 397                 name = strchr(typelist[i]->tp_name, '.');
 398                 assert (name != NULL);
 399                 Py_INCREF(typelist[i]);
 400                 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
 401         }
 402 }