Objects/iterobject.c

   1 /* Iterator objects */
   2
   3 #include "Python.h"
   4
   5 typedef struct {
   6         PyObject_HEAD
   7         long      it_index;
   8         PyObject *it_seq; /* Set to NULL when iterator is exhausted */
   9 } seqiterobject;
  10
  11 PyObject *
  12 PySeqIter_New(PyObject *seq)
  13 {
  14         seqiterobject *it;
  15
  16         if (!PySequence_Check(seq)) {
  17                 PyErr_BadInternalCall();
  18                 return NULL;
  19         }
  20         it = PyObject_GC_New(seqiterobject, &PySeqIter_Type);
  21         if (it == NULL)
  22                 return NULL;
  23         it->it_index = 0;
  24         Py_INCREF(seq);
  25         it->it_seq = seq;
  26         _PyObject_GC_TRACK(it);
  27         return (PyObject *)it;
  28 }
  29
  30 static void
  31 iter_dealloc(seqiterobject *it)
  32 {
  33         _PyObject_GC_UNTRACK(it);
  34         Py_XDECREF(it->it_seq);
  35         PyObject_GC_Del(it);
  36 }
  37
  38 static int
  39 iter_traverse(seqiterobject *it, visitproc visit, void *arg)
  40 {
  41         Py_VISIT(it->it_seq);
  42         return 0;
  43 }
  44
  45 static PyObject *
  46 iter_iternext(PyObject *iterator)
  47 {
  48         seqiterobject *it;
  49         PyObject *seq;
  50         PyObject *result;
  51
  52         assert(PySeqIter_Check(iterator));
  53         it = (seqiterobject *)iterator;
  54         seq = it->it_seq;
  55         if (seq == NULL)
  56                 return NULL;
  57
  58         result = PySequence_GetItem(seq, it->it_index);
  59         if (result != NULL) {
  60                 it->it_index++;
  61                 return result;
  62         }
  63         if (PyErr_ExceptionMatches(PyExc_IndexError) ||
  64             PyErr_ExceptionMatches(PyExc_StopIteration))
  65         {
  66                 PyErr_Clear();
  67                 Py_DECREF(seq);
  68                 it->it_seq = NULL;
  69         }
  70         return NULL;
  71 }
  72
  73 static PyObject *
  74 iter_len(seqiterobject *it)
  75 {
  76         Py_ssize_t seqsize, len;
  77
  78         if (it->it_seq) {
  79                 seqsize = PySequence_Size(it->it_seq);
  80                 if (seqsize == -1)
  81                         return NULL;
  82                 len = seqsize - it->it_index;
  83                 if (len >= 0)
  84                         return PyInt_FromSsize_t(len);
  85         }
  86         return PyInt_FromLong(0);
  87 }
  88
  89 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
  90
  91 static PyMethodDef seqiter_methods[] = {
  92         {"__length_hint__", (PyCFunction)iter_len, METH_NOARGS, length_hint_doc},
  93         {NULL,          NULL}           /* sentinel */
  94 };
  95
  96 PyTypeObject PySeqIter_Type = {
  97         PyObject_HEAD_INIT(&PyType_Type)
  98         0,                                      /* ob_size */
  99         "iterator",                             /* tp_name */
 100         sizeof(seqiterobject),                  /* tp_basicsize */
 101         0,                                      /* tp_itemsize */
 102         /* methods */
 103         (destructor)iter_dealloc,               /* tp_dealloc */
 104         0,                                      /* tp_print */
 105         0,                                      /* tp_getattr */
 106         0,                                      /* tp_setattr */
 107         0,                                      /* tp_compare */
 108         0,                                      /* tp_repr */
 109         0,                                      /* tp_as_number */
 110         0,                                      /* tp_as_sequence */
 111         0,                                      /* tp_as_mapping */
 112         0,                                      /* tp_hash */
 113         0,                                      /* tp_call */
 114         0,                                      /* tp_str */
 115         PyObject_GenericGetAttr,                /* tp_getattro */
 116         0,                                      /* tp_setattro */
 117         0,                                      /* tp_as_buffer */
 118         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
 119         0,                                      /* tp_doc */
 120         (traverseproc)iter_traverse,            /* tp_traverse */
 121         0,                                      /* tp_clear */
 122         0,                                      /* tp_richcompare */
 123         0,                                      /* tp_weaklistoffset */
 124         PyObject_SelfIter,                      /* tp_iter */
 125         iter_iternext,                          /* tp_iternext */
 126         seqiter_methods,                        /* tp_methods */
 127         0,                                      /* tp_members */
 128 };
 129
 130 /* -------------------------------------- */
 131
 132 typedef struct {
 133         PyObject_HEAD
 134         PyObject *it_callable; /* Set to NULL when iterator is exhausted */
 135         PyObject *it_sentinel; /* Set to NULL when iterator is exhausted */
 136 } calliterobject;
 137
 138 PyObject *
 139 PyCallIter_New(PyObject *callable, PyObject *sentinel)
 140 {
 141         calliterobject *it;
 142         it = PyObject_GC_New(calliterobject, &PyCallIter_Type);
 143         if (it == NULL)
 144                 return NULL;
 145         Py_INCREF(callable);
 146         it->it_callable = callable;
 147         Py_INCREF(sentinel);
 148         it->it_sentinel = sentinel;
 149         _PyObject_GC_TRACK(it);
 150         return (PyObject *)it;
 151 }
 152 static void
 153 calliter_dealloc(calliterobject *it)
 154 {
 155         _PyObject_GC_UNTRACK(it);
 156         Py_XDECREF(it->it_callable);
 157         Py_XDECREF(it->it_sentinel);
 158         PyObject_GC_Del(it);
 159 }
 160
 161 static int
 162 calliter_traverse(calliterobject *it, visitproc visit, void *arg)
 163 {
 164         Py_VISIT(it->it_callable);
 165         Py_VISIT(it->it_sentinel);
 166         return 0;
 167 }
 168
 169 static PyObject *
 170 calliter_iternext(calliterobject *it)
 171 {
 172         if (it->it_callable != NULL) {
 173                 PyObject *args = PyTuple_New(0);
 174                 PyObject *result;
 175                 if (args == NULL)
 176                         return NULL;
 177                 result = PyObject_Call(it->it_callable, args, NULL);
 178                 Py_DECREF(args);
 179                 if (result != NULL) {
 180                         int ok;
 181                         ok = PyObject_RichCompareBool(result,
 182                                                       it->it_sentinel,
 183                                                       Py_EQ);
 184                         if (ok == 0)
 185                                 return result; /* Common case, fast path */
 186                         Py_DECREF(result);
 187                         if (ok > 0) {
 188                                 Py_CLEAR(it->it_callable);
 189                                 Py_CLEAR(it->it_sentinel);
 190                         }
 191                 }
 192                 else if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
 193                         PyErr_Clear();
 194                         Py_CLEAR(it->it_callable);
 195                         Py_CLEAR(it->it_sentinel);
 196                 }
 197         }
 198         return NULL;
 199 }
 200
 201 PyTypeObject PyCallIter_Type = {
 202         PyObject_HEAD_INIT(&PyType_Type)
 203         0,                                      /* ob_size */
 204         "callable-iterator",                    /* tp_name */
 205         sizeof(calliterobject),                 /* tp_basicsize */
 206         0,                                      /* tp_itemsize */
 207         /* methods */
 208         (destructor)calliter_dealloc,           /* tp_dealloc */
 209         0,                                      /* tp_print */
 210         0,                                      /* tp_getattr */
 211         0,                                      /* tp_setattr */
 212         0,                                      /* tp_compare */
 213         0,                                      /* tp_repr */
 214         0,                                      /* tp_as_number */
 215         0,                                      /* tp_as_sequence */
 216         0,                                      /* tp_as_mapping */
 217         0,                                      /* tp_hash */
 218         0,                                      /* tp_call */
 219         0,                                      /* tp_str */
 220         PyObject_GenericGetAttr,                /* tp_getattro */
 221         0,                                      /* tp_setattro */
 222         0,                                      /* tp_as_buffer */
 223         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
 224         0,                                      /* tp_doc */
 225         (traverseproc)calliter_traverse,        /* tp_traverse */
 226         0,                                      /* tp_clear */
 227         0,                                      /* tp_richcompare */
 228         0,                                      /* tp_weaklistoffset */
 229         PyObject_SelfIter,                      /* tp_iter */
 230         (iternextfunc)calliter_iternext,        /* tp_iternext */
 231         0,                                      /* tp_methods */
 232 };