How to create a generator/iterator with the Python C API?

Solution 1:

Below is a simple implementation of module spam with one function myiter(int) returning iterator:

import spam
for i in spam.myiter(10):
    print i

prints numbers from 0 to 9.

It is simpler then your case but shows main points: defining object with standard __iter__() and next() methods, and implementing iterator behaviour including raising StopIteration when appropriate.

In your case iterator object needs to hold reference to Sequence (so you'll need deallocator method for it to Py_DECREF it). The sequence itself needs to implement __iter()__ and create an iterator inside it.

Structure containing state of iterator. (In your version instead of m, it would have reference to Sequence.)

typedef struct {
  PyObject_HEAD
  long int m;
  long int i;
} spam_MyIter;

Iterator's __iter__() method. It always simply returns self. It allows for both iterator and collection to be treated the same in constructs like for ... in ....

PyObject* spam_MyIter_iter(PyObject *self)
{
  Py_INCREF(self);
  return self;
}

Implementation of our iteration: next() method.

PyObject* spam_MyIter_iternext(PyObject *self)
{
  spam_MyIter *p = (spam_MyIter *)self;
  if (p->i < p->m) {
    PyObject *tmp = Py_BuildValue("l", p->i);
    (p->i)++;
    return tmp;
  } else {
    /* Raising of standard StopIteration exception with empty value. */
    PyErr_SetNone(PyExc_StopIteration);
    return NULL;
  }
}

We need extended version of PyTypeObject structure to provide Python with information about __iter__() and next(). We want them to be called efficiently, so no name-based lookup in dictionary.

static PyTypeObject spam_MyIterType = {
    PyObject_HEAD_INIT(NULL)
    0,                         /*ob_size*/
    "spam._MyIter",            /*tp_name*/
    sizeof(spam_MyIter),       /*tp_basicsize*/
    0,                         /*tp_itemsize*/
    0,                         /*tp_dealloc*/
    0,                         /*tp_print*/
    0,                         /*tp_getattr*/
    0,                         /*tp_setattr*/
    0,                         /*tp_compare*/
    0,                         /*tp_repr*/
    0,                         /*tp_as_number*/
    0,                         /*tp_as_sequence*/
    0,                         /*tp_as_mapping*/
    0,                         /*tp_hash */
    0,                         /*tp_call*/
    0,                         /*tp_str*/
    0,                         /*tp_getattro*/
    0,                         /*tp_setattro*/
    0,                         /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_ITER,
      /* tp_flags: Py_TPFLAGS_HAVE_ITER tells python to
         use tp_iter and tp_iternext fields. */
    "Internal myiter iterator object.",           /* tp_doc */
    0,  /* tp_traverse */
    0,  /* tp_clear */
    0,  /* tp_richcompare */
    0,  /* tp_weaklistoffset */
    spam_MyIter_iter,  /* tp_iter: __iter__() method */
    spam_MyIter_iternext  /* tp_iternext: next() method */
};

myiter(int) function creates iterator.

static PyObject *
spam_myiter(PyObject *self, PyObject *args)
{
  long int m;
  spam_MyIter *p;

  if (!PyArg_ParseTuple(args, "l", &m))  return NULL;

  /* I don't need python callable __init__() method for this iterator,
     so I'll simply allocate it as PyObject and initialize it by hand. */

  p = PyObject_New(spam_MyIter, &spam_MyIterType);
  if (!p) return NULL;

  /* I'm not sure if it's strictly necessary. */
  if (!PyObject_Init((PyObject *)p, &spam_MyIterType)) {
    Py_DECREF(p);
    return NULL;
  }

  p->m = m;
  p->i = 0;
  return (PyObject *)p;
}

The rest is pretty boring...

static PyMethodDef SpamMethods[] = {
    {"myiter",  spam_myiter, METH_VARARGS, "Iterate from i=0 while i<m."},
    {NULL, NULL, 0, NULL}        /* Sentinel */
};

PyMODINIT_FUNC
initspam(void)
{
  PyObject* m;

  spam_MyIterType.tp_new = PyType_GenericNew;
  if (PyType_Ready(&spam_MyIterType) < 0)  return;

  m = Py_InitModule("spam", SpamMethods);

  Py_INCREF(&spam_MyIterType);
  PyModule_AddObject(m, "_MyIter", (PyObject *)&spam_MyIterType);
}

Solution 2:

In Sequence_data, you must either return a new PyInt instance or throw a StopIteration exception which tells the code outside that there are no more values. See PEP 255 for details and 9.10 Generators.

See Iterator Protocol for helper functions in the Python/C API.