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 static PyTypeObject partial_type = {
 200         PyVarObject_HEAD_INIT(NULL, 0)
 201         "functools.partial",            /* tp_name */
 202         sizeof(partialobject),          /* tp_basicsize */
 203         0,                              /* tp_itemsize */
 204         /* methods */
 205         (destructor)partial_dealloc,    /* tp_dealloc */
 206         0,                              /* tp_print */
 207         0,                              /* tp_getattr */
 208         0,                              /* tp_setattr */
 209         0,                              /* tp_compare */
 210         0,                              /* tp_repr */
 211         0,                              /* tp_as_number */
 212         0,                              /* tp_as_sequence */
 213         0,                              /* tp_as_mapping */
 214         0,                              /* tp_hash */
 215         (ternaryfunc)partial_call,      /* tp_call */
 216         0,                              /* tp_str */
 217         PyObject_GenericGetAttr,        /* tp_getattro */
 218         PyObject_GenericSetAttr,        /* tp_setattro */
 219         0,                              /* tp_as_buffer */
 220         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 221                 Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_WEAKREFS, /* tp_flags */
 222         partial_doc,                    /* tp_doc */
 223         (traverseproc)partial_traverse, /* tp_traverse */
 224         0,                              /* tp_clear */
 225         0,                              /* tp_richcompare */
 226         offsetof(partialobject, weakreflist),   /* tp_weaklistoffset */
 227         0,                              /* tp_iter */
 228         0,                              /* tp_iternext */
 229         0,                              /* tp_methods */
 230         partial_memberlist,             /* tp_members */
 231         partial_getsetlist,             /* tp_getset */
 232         0,                              /* tp_base */
 233         0,                              /* tp_dict */
 234         0,                              /* tp_descr_get */
 235         0,                              /* tp_descr_set */
 236         offsetof(partialobject, dict),  /* tp_dictoffset */
 237         0,                              /* tp_init */
 238         0,                              /* tp_alloc */
 239         partial_new,                    /* tp_new */
 240         PyObject_GC_Del,                /* tp_free */
 241 };
 242
 243
 244 /* module level code ********************************************************/
 245
 246 PyDoc_STRVAR(module_doc,
 247 "Tools that operate on functions.");
 248
 249 static PyMethodDef module_methods[] = {
 250         {NULL,          NULL}           /* sentinel */
 251 };
 252
 253 PyMODINIT_FUNC
 254 init_functools(void)
 255 {
 256         int i;
 257         PyObject *m;
 258         char *name;
 259         PyTypeObject *typelist[] = {
 260                 &partial_type,
 261                 NULL
 262         };
 263
 264         m = Py_InitModule3("_functools", module_methods, module_doc);
 265         if (m == NULL)
 266                 return;
 267
 268         for (i=0 ; typelist[i] != NULL ; i++) {
 269                 if (PyType_Ready(typelist[i]) < 0)
 270                         return;
 271                 name = strchr(typelist[i]->tp_name, '.');
 272                 assert (name != NULL);
 273                 Py_INCREF(typelist[i]);
 274                 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
 275         }
 276 }