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