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 } enumobject;
  10
  11 PyTypeObject PyEnum_Type;
  12
  13 static PyObject *
  14 enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  15 {
  16         enumobject *en;
  17         PyObject *seq = NULL;
  18         static char *kwlist[] = {"sequence", 0};
  19
  20         if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:enumerate", kwlist,
  21                                          &seq))
  22                 return NULL;
  23
  24         en = (enumobject *)type->tp_alloc(type, 0);
  25         if (en == NULL)
  26                 return NULL;
  27         en->en_index = 0;
  28         en->en_sit = PyObject_GetIter(seq);
  29         if (en->en_sit == NULL) {
  30                 Py_DECREF(en);
  31                 return NULL;
  32         }
  33         return (PyObject *)en;
  34 }
  35
  36 static void
  37 enum_dealloc(enumobject *en)
  38 {
  39         PyObject_GC_UnTrack(en);
  40         Py_XDECREF(en->en_sit);
  41         en->ob_type->tp_free(en);
  42 }
  43
  44 static int
  45 enum_traverse(enumobject *en, visitproc visit, void *arg)
  46 {
  47         if (en->en_sit)
  48                 return visit(en->en_sit, arg);
  49         return 0;
  50 }
  51
  52 static PyObject *
  53 enum_next(enumobject *en)
  54 {
  55         PyObject *result;
  56         PyObject *next_index;
  57         PyObject *next_item;
  58
  59         result = PyTuple_New(2);
  60         if (result == NULL)
  61                 return NULL;
  62
  63         next_index = PyInt_FromLong(en->en_index);
  64         if (next_index == NULL) {
  65                 Py_DECREF(result);
  66                 return NULL;
  67         }
  68         PyTuple_SET_ITEM(result, 0, next_index);
  69
  70         next_item = PyIter_Next(en->en_sit);
  71         if (next_item == NULL) {
  72                 Py_DECREF(result);
  73                 return NULL;
  74         }
  75
  76         en->en_index++;
  77         PyTuple_SET_ITEM(result, 1, next_item);
  78         return result;
  79 }
  80
  81 static PyObject *
  82 enum_getiter(PyObject *en)
  83 {
  84         Py_INCREF(en);
  85         return en;
  86 }
  87
  88 PyDoc_STRVAR(enum_doc,
  89 "enumerate(iterable) -> create an enumerating-iterator");
  90
  91 PyTypeObject PyEnum_Type = {
  92         PyObject_HEAD_INIT(&PyType_Type)
  93         0,                              /* ob_size */
  94         "enumerate",                    /* tp_name */
  95         sizeof(enumobject),             /* tp_basicsize */
  96         0,                              /* tp_itemsize */
  97         /* methods */
  98         (destructor)enum_dealloc,       /* tp_dealloc */
  99         0,                              /* tp_print */
 100         0,                              /* tp_getattr */
 101         0,                              /* tp_setattr */
 102         0,                              /* tp_compare */
 103         0,                              /* tp_repr */
 104         0,                              /* tp_as_number */
 105         0,                              /* tp_as_sequence */
 106         0,                              /* tp_as_mapping */
 107         0,                              /* tp_hash */
 108         0,                              /* tp_call */
 109         0,                              /* tp_str */
 110         PyObject_GenericGetAttr,        /* tp_getattro */
 111         0,                              /* tp_setattro */
 112         0,                              /* tp_as_buffer */
 113         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
 114                 Py_TPFLAGS_BASETYPE,    /* tp_flags */
 115         enum_doc,                       /* tp_doc */
 116         (traverseproc)enum_traverse,    /* tp_traverse */
 117         0,                              /* tp_clear */
 118         0,                              /* tp_richcompare */
 119         0,                              /* tp_weaklistoffset */
 120         (getiterfunc)enum_getiter,      /* tp_iter */
 121         (iternextfunc)enum_next,        /* tp_iternext */
 122         0,                              /* tp_methods */
 123         0,                              /* tp_members */
 124         0,                              /* tp_getset */
 125         0,                              /* tp_base */
 126         0,                              /* tp_dict */
 127         0,                              /* tp_descr_get */
 128         0,                              /* tp_descr_set */
 129         0,                              /* tp_dictoffset */
 130         0,                              /* tp_init */
 131         PyType_GenericAlloc,            /* tp_alloc */
 132         enum_new,                       /* tp_new */
 133         PyObject_GC_Del,                /* tp_free */
 134 };