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         PyVarObject_HEAD_INIT(&PyType_Type, 0)
  98         "iterator",                             /* tp_name */
  99         sizeof(seqiterobject),                  /* tp_basicsize */
 100         0,                                      /* tp_itemsize */
 101         /* methods */
 102         (destructor)iter_dealloc,               /* tp_dealloc */
 103         0,                                      /* tp_print */
 104         0,                                      /* tp_getattr */
 105         0,                                      /* tp_setattr */
 106         0,                                      /* tp_compare */
 107         0,                                      /* tp_repr */
 108         0,                                      /* tp_as_number */
 109         0,                                      /* tp_as_sequence */
 110         0,                                      /* tp_as_mapping */
 111         0,                                      /* tp_hash */
 112         0,                                      /* tp_call */
 113         0,                                      /* tp_str */
 114         PyObject_GenericGetAttr,                /* tp_getattro */
 115         0,                                      /* tp_setattro */
 116         0,                                      /* tp_as_buffer */
 117         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
 118         0,                                      /* tp_doc */
 119         (traverseproc)iter_traverse,            /* tp_traverse */
 120         0,                                      /* tp_clear */
 121         0,                                      /* tp_richcompare */
 122         0,                                      /* tp_weaklistoffset */
 123         PyObject_SelfIter,                      /* tp_iter */
 124         iter_iternext,                          /* tp_iternext */
 125         seqiter_methods,                        /* tp_methods */
 126         0,                                      /* tp_members */
 127 };
 128
 129 /* -------------------------------------- */
 130
 131 typedef struct {
 132         PyObject_HEAD
 133         PyObject *it_callable; /* Set to NULL when iterator is exhausted */
 134         PyObject *it_sentinel; /* Set to NULL when iterator is exhausted */
 135 } calliterobject;
 136
 137 PyObject *
 138 PyCallIter_New(PyObject *callable, PyObject *sentinel)
 139 {
 140         calliterobject *it;
 141         it = PyObject_GC_New(calliterobject, &PyCallIter_Type);
 142         if (it == NULL)
 143                 return NULL;
 144         Py_INCREF(callable);
 145         it->it_callable = callable;
 146         Py_INCREF(sentinel);
 147         it->it_sentinel = sentinel;
 148         _PyObject_GC_TRACK(it);
 149         return (PyObject *)it;
 150 }
 151 static void
 152 calliter_dealloc(calliterobject *it)
 153 {
 154         _PyObject_GC_UNTRACK(it);
 155         Py_XDECREF(it->it_callable);
 156         Py_XDECREF(it->it_sentinel);
 157         PyObject_GC_Del(it);
 158 }
 159
 160 static int
 161 calliter_traverse(calliterobject *it, visitproc visit, void *arg)
 162 {
 163         Py_VISIT(it->it_callable);
 164         Py_VISIT(it->it_sentinel);
 165         return 0;
 166 }
 167
 168 static PyObject *
 169 calliter_iternext(calliterobject *it)
 170 {
 171         if (it->it_callable != NULL) {
 172                 PyObject *args = PyTuple_New(0);
 173                 PyObject *result;
 174                 if (args == NULL)
 175                         return NULL;
 176                 result = PyObject_Call(it->it_callable, args, NULL);
 177                 Py_DECREF(args);
 178                 if (result != NULL) {
 179                         int ok;
 180                         ok = PyObject_RichCompareBool(result,
 181                                                       it->it_sentinel,
 182                                                       Py_EQ);
 183                         if (ok == 0)
 184                                 return result; /* Common case, fast path */
 185                         Py_DECREF(result);
 186                         if (ok > 0) {
 187                                 Py_CLEAR(it->it_callable);
 188                                 Py_CLEAR(it->it_sentinel);
 189                         }
 190                 }
 191                 else if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
 192                         PyErr_Clear();
 193                         Py_CLEAR(it->it_callable);
 194                         Py_CLEAR(it->it_sentinel);
 195                 }
 196         }
 197         return NULL;
 198 }
 199
 200 PyTypeObject PyCallIter_Type = {
 201         PyVarObject_HEAD_INIT(&PyType_Type, 0)
 202         "callable-iterator",                    /* tp_name */
 203         sizeof(calliterobject),                 /* tp_basicsize */
 204         0,                                      /* tp_itemsize */
 205         /* methods */
 206         (destructor)calliter_dealloc,           /* tp_dealloc */
 207         0,                                      /* tp_print */
 208         0,                                      /* tp_getattr */
 209         0,                                      /* tp_setattr */
 210         0,                                      /* tp_compare */
 211         0,                                      /* tp_repr */
 212         0,                                      /* tp_as_number */
 213         0,                                      /* tp_as_sequence */
 214         0,                                      /* tp_as_mapping */
 215         0,                                      /* tp_hash */
 216         0,                                      /* tp_call */
 217         0,                                      /* tp_str */
 218         PyObject_GenericGetAttr,                /* tp_getattro */
 219         0,                                      /* tp_setattro */
 220         0,                                      /* tp_as_buffer */
 221         Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
 222         0,                                      /* tp_doc */
 223         (traverseproc)calliter_traverse,        /* tp_traverse */
 224         0,                                      /* tp_clear */
 225         0,                                      /* tp_richcompare */
 226         0,                                      /* tp_weaklistoffset */
 227         PyObject_SelfIter,                      /* tp_iter */
 228         (iternextfunc)calliter_iternext,        /* tp_iternext */
 229         0,                                      /* tp_methods */
 230 };