Cython close to hand-crafted C code for generators

I did a couple of experiments compiling itertools with the new generator support in Cython. In CPython, the itertools module is actually written in hand tuned C and does very little computation in its generators, so I knew it would be hard to reach with generated code. But Cython does a pretty good job.

Something as trivial as chain() is exactly as fast as in the C implementation, but compared to the more than 60 lines of C code, it is certainly a lot more readable in Cython:

def chain(*iterables):
    """Make an iterator that returns elements from the first iterable
    until it is exhausted, then proceeds to the next iterable, until
    all of the iterables are exhausted. Used for treating consecutive
    sequences as a single sequence.
    """
    for it in iterables:
        for element in it:
            yield element

Other functions, like islice(), are faster in C, partly because CPython actually takes a couple of shortcuts, e.g. by only looking up the iterator slot method once. You cannot do that in Python code, and I wanted to keep the implementation compatible with regular Python. Specifically, the C speed advantage for islice() is currently about 30-50% in general, although the Cython implementation can also be up to 10% faster for some cases, e.g. when extracting only a couple of items from the middle of a longer sequence. The C implementation is about 90 lines, here is the Cython implementation:

import sys

import cython

# Python 2/3 compatibility

_max_size = cython.declare(cython.Py_ssize_t,
        getattr(sys, "maxsize", getattr(sys, "maxint", None)))

@cython.locals(i=cython.Py_ssize_t, nexti=cython.Py_ssize_t,
               start=cython.Py_ssize_t, stop=cython.Py_ssize_t, step=cython.Py_ssize_t)
def islice(iterable, *args):
    """Make an iterator that returns selected elements from the
    iterable. If start is non-zero, then elements from the iterable
    are skipped until start is reached. Afterward, elements are
    returned consecutively unless step is set higher than one which
    results in items being skipped. If stop is None, then iteration
    continues until the iterator is exhausted, if at all; otherwise,
    it stops at the specified position. Unlike regular slicing,
    islice() does not support negative values for start, stop, or
    step. Can be used to extract related fields from data where the
    internal structure has been flattened (for example, a multi-line
    report may list a name field on every third line).
    """
    s = slice(*args)
    start = s.start or 0
    stop = s.stop or _max_size
    step = s.step or 1
    if start = stop:
        return
    nexti = start
    for i, element in enumerate(iterable):
        if i == nexti:
            yield element
            nexti += step
            if nexti >= stop or nexti 

Here is one that is conceptually quite simple: count(). I had to optimise it quite a bit, because the iteration code in the C code is extremely tight. Even the tuned version below runs about 10% slower than the hand tuned C version, which is about 230 lines long.


@cython.locals(i=cython.Py_ssize_t)

def count(n=0):
    """Make an iterator that returns consecutive integers starting
    with n. If not specified n defaults to zero. Often used as an argument to imap()
    to generate consecutive data points. Also, used with zip() to add
    sequence numbers.
    """
    try:
        i = n
    except OverflowError:
        i = _max_size # skip i-loop
    else:
        n = _max_size # first value after i-loop
    while i 

Note that all of the above generators execute in the order of microseconds, so even a slow-down of 50% will likely not be measurable in real world code.

So far, I did not try any of the more fancy functions in itertools (those that actually do something). The Cython project has announced a Google Summer of Code project with exactly the intent to rewrite some of the C stdlib modules of CPython in pure Python code with Cython compiler hints. So I leave this exercise to interested readers for now.