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