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