Objects/enumobject.c

   1 /* enumerate object */
   2
   3 #include "Python.h"
   4
   5 typedef struct {
   6         PyObject_HEAD
   7         long      en_index;        /* current index of enumeration */
   8         PyObject* en_sit;          /* secondary iterator of enumeration */
   9         PyObject* en_result;       /* result tuple  */
  10 } enumobject;
  11
  12 static PyObject *
  13 enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  14 {
  15         enumobject *en;
  16         PyObject *seq = NULL;
  17         static char *kwlist[] = {"sequence", 0};
  18
  19         if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:enumerate", kwlist,
  20                                          &seq))
  21                 return NULL;
  22
  23         en = (enumobject *)type->tp_alloc(type, 0);
  24         if (en == NULL)
  25                 return NULL;
  26         en->en_index = 0;
  27         en->en_sit = PyObject_GetIter(seq);
  28         if (en->en_sit == NULL) {
  29                 Py_DECREF(en);
  30                 return NULL;
  31         }
  32         en->en_result = PyTuple_Pack(2, Py_None, Py_None);
  33         if (en->en_result == NULL) {
  34                 Py_DECREF(en);
  35                 return NULL;
  36         }
  37         return (PyObject *)en;
  38 }
  39
  40 static void
  41 enum_dealloc(enumobject *en)
  42 {
  43         PyObject_GC_UnTrack(en);
  44         Py_XDECREF(en->en_sit);
  45         Py_XDECREF(en->en_result);
  46         en->ob_type->tp_free(en);
  47 }
  48
  49 static int
  50 enum_traverse(enumobject *en, visitproc visit, void *arg)
  51 {
  52         Py_VISIT(en->en_sit);
  53         Py_VISIT(en->en_result);
  54         return 0;
  55 }
  56
  57 static PyObject *
  58 enum_next(enumobject *en)
  59 {
  60         PyObject *next_index;
  61         PyObject *next_item;
  62         PyObject *result = en->en_result;
  63         PyObject *it = en->en_sit;
  64
  65         next_item = (*it->ob_type->tp_iternext)(it);
  66         if (next_item == NULL)
  67                 return NULL;
  68
  69         next_index = PyInt_FromLong(en->en_index);
  70         if (next_index == NULL) {
  71                 Py_DECREF(next_item);
  72                 return NULL;
  73         }
  74         en->en_index++;
  75
  76         if (result->ob_refcnt == 1) {
  77                 Py_INCREF(result);
  78                 Py_DECREF(PyTuple_GET_ITEM(result, 0));
  79                 Py_DECREF(PyTuple_GET_ITEM(result, 1));
  80         } else {
  81                 result = PyTuple_New(2);
  82                 if (result == NULL) {
  83                         Py_DECREF(next_index);
  84                         Py_DECREF(next_item);
  85                         return NULL;
  86                 }
  87         }
  88         PyTuple_SET_ITEM(result, 0, next_index);
  89         PyTuple_SET_ITEM(result, 1, next_item);
  90         return result;
  91 }
  92
  93 PyDoc_STRVAR(enum_doc,
  94 "enumerate(iterable) -> iterator for index, value of iterable\n"
  95 "\n"
  96 "Return an enumerate object.  iterable must be an other object that supports\n"
  97 "iteration.  The enumerate object yields pairs containing a count (from\n"
  98 "zero) and a value yielded by the iterable argument.  enumerate is useful\n"
  99 "for obtaining an indexed list: (0, seq[0]), (1, seq[1]), (2, seq[2]), ...");
 100
 101 PyTypeObject PyEnum_Type = {
 102         PyObject_HEAD_INIT(&PyType_Type)
 103         0,                              /* ob_size */
 104         "enumerate",                    /* tp_name */
 105         sizeof(enumobject),             /* tp_basicsize */
 106         0,                              /* tp_itemsize */
 107         /* methods */
 108         (destructor)enum_dealloc,       /* tp_dealloc */
 109         0,                              /* tp_print */
 110         0,                              /* tp_getattr */
 111         0,                              /* tp_setattr */
 112         0,                              /* tp_compare */
 113         0,                              /* tp_repr */
 114         0,                              /* tp_as_number */
 115         0,                              /* tp_as_sequence */
 116         0,                              /* tp_as_mapping */
 117         0,                              /* tp_hash */
 118         0,                              /* tp_call */
 119         0,                              /* tp_str */
 120         PyObject_GenericGetAttr,        /* tp_getattro */
 121         0,                              /* tp_setattro */
 122         0,                              /* tp_as_buffer */
 123         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 124                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 125         enum_doc,                       /* tp_doc */
 126         (traverseproc)enum_traverse,    /* tp_traverse */
 127         0,                              /* tp_clear */
 128         0,                              /* tp_richcompare */
 129         0,                              /* tp_weaklistoffset */
 130         PyObject_SelfIter,              /* tp_iter */
 131         (iternextfunc)enum_next,        /* tp_iternext */
 132         0,                              /* tp_methods */
 133         0,                              /* tp_members */
 134         0,                              /* tp_getset */
 135         0,                              /* tp_base */
 136         0,                              /* tp_dict */
 137         0,                              /* tp_descr_get */
 138         0,                              /* tp_descr_set */
 139         0,                              /* tp_dictoffset */
 140         0,                              /* tp_init */
 141         PyType_GenericAlloc,            /* tp_alloc */
 142         enum_new,                       /* tp_new */
 143         PyObject_GC_Del,                /* tp_free */
 144 };
 145
 146 /* Reversed Object ***************************************************************/
 147
 148 typedef struct {
 149         PyObject_HEAD
 150         Py_ssize_t      index;
 151         PyObject* seq;
 152 } reversedobject;
 153
 154 static PyObject *
 155 reversed_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 156 {
 157         Py_ssize_t n;
 158         PyObject *seq;
 159         reversedobject *ro;
 160
 161         if (!PyArg_UnpackTuple(args, "reversed", 1, 1, &seq))
 162                 return NULL;
 163
 164         if (PyObject_HasAttrString(seq, "__reversed__"))
 165                 return PyObject_CallMethod(seq, "__reversed__", NULL);
 166
 167         if (!PySequence_Check(seq)) {
 168                 PyErr_SetString(PyExc_TypeError,
 169                                 "argument to reversed() must be a sequence");
 170                 return NULL;
 171         }
 172
 173         n = PySequence_Size(seq);
 174         if (n == -1)
 175                 return NULL;
 176
 177         ro = (reversedobject *)type->tp_alloc(type, 0);
 178         if (ro == NULL)
 179                 return NULL;
 180
 181         ro->index = n-1;
 182         Py_INCREF(seq);
 183         ro->seq = seq;
 184         return (PyObject *)ro;
 185 }
 186
 187 static void
 188 reversed_dealloc(reversedobject *ro)
 189 {
 190         PyObject_GC_UnTrack(ro);
 191         Py_XDECREF(ro->seq);
 192         ro->ob_type->tp_free(ro);
 193 }
 194
 195 static int
 196 reversed_traverse(reversedobject *ro, visitproc visit, void *arg)
 197 {
 198         Py_VISIT(ro->seq);
 199         return 0;
 200 }
 201
 202 static PyObject *
 203 reversed_next(reversedobject *ro)
 204 {
 205         PyObject *item;
 206         Py_ssize_t index = ro->index;
 207
 208         if (index >= 0) {
 209                 item = PySequence_GetItem(ro->seq, index);
 210                 if (item != NULL) {
 211                         ro->index--;
 212                         return item;
 213                 }
 214                 if (PyErr_ExceptionMatches(PyExc_IndexError) ||
 215                     PyErr_ExceptionMatches(PyExc_StopIteration))
 216                         PyErr_Clear();
 217         }
 218         ro->index = -1;
 219         Py_CLEAR(ro->seq);
 220         return NULL;
 221 }
 222
 223 PyDoc_STRVAR(reversed_doc,
 224 "reversed(sequence) -> reverse iterator over values of the sequence\n"
 225 "\n"
 226 "Return a reverse iterator");
 227
 228 static PyObject *
 229 reversed_len(reversedobject *ro)
 230 {
 231         Py_ssize_t position, seqsize;
 232
 233         if (ro->seq == NULL)
 234                 return PyInt_FromLong(0);
 235         seqsize = PySequence_Size(ro->seq);
 236         if (seqsize == -1)
 237                 return NULL;
 238         position = ro->index + 1;
 239         return PyInt_FromSsize_t((seqsize < position)  ?  0  :  position);
 240 }
 241
 242 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
 243
 244 static PyMethodDef reversediter_methods[] = {
 245         {"__length_hint__", (PyCFunction)reversed_len, METH_NOARGS, length_hint_doc},
 246         {NULL,          NULL}           /* sentinel */
 247 };
 248
 249 PyTypeObject PyReversed_Type = {
 250         PyObject_HEAD_INIT(&PyType_Type)
 251         0,                              /* ob_size */
 252         "reversed",                     /* tp_name */
 253         sizeof(reversedobject),         /* tp_basicsize */
 254         0,                              /* tp_itemsize */
 255         /* methods */
 256         (destructor)reversed_dealloc,   /* tp_dealloc */
 257         0,                              /* tp_print */
 258         0,                              /* tp_getattr */
 259         0,                              /* tp_setattr */
 260         0,                              /* tp_compare */
 261         0,                              /* tp_repr */
 262         0,                              /* tp_as_number */
 263         0,                              /* tp_as_sequence */
 264         0,                              /* tp_as_mapping */
 265         0,                              /* tp_hash */
 266         0,                              /* tp_call */
 267         0,                              /* tp_str */
 268         PyObject_GenericGetAttr,        /* tp_getattro */
 269         0,                              /* tp_setattro */
 270         0,                              /* tp_as_buffer */
 271         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 272                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 273         reversed_doc,                   /* tp_doc */
 274         (traverseproc)reversed_traverse,/* tp_traverse */
 275         0,                              /* tp_clear */
 276         0,                              /* tp_richcompare */
 277         0,                              /* tp_weaklistoffset */
 278         PyObject_SelfIter,              /* tp_iter */
 279         (iternextfunc)reversed_next,    /* tp_iternext */
 280         reversediter_methods,           /* tp_methods */
 281         0,                              /* tp_members */
 282         0,                              /* tp_getset */
 283         0,                              /* tp_base */
 284         0,                              /* tp_dict */
 285         0,                              /* tp_descr_get */
 286         0,                              /* tp_descr_set */
 287         0,                              /* tp_dictoffset */
 288         0,                              /* tp_init */
 289         PyType_GenericAlloc,            /* tp_alloc */
 290         reversed_new,                   /* tp_new */
 291         PyObject_GC_Del,                /* tp_free */
 292 };