Retrieving the nth term of an infinite stream

I recently learned how generators can be infinite in Python. For example, the infinite sequence can be the triangular numbers:
$$1, 3, 6, 10,...$$

def triangle_stream():
    """ Infinite stream of triangle numbers. """
    x = step = 1
    yield x
    while True:
        step += 1
        x += step
        yield x

I wrote a function that returns the nth term in an infinite sequence.

def nth_term(generator, n):
    """ Returns the nth term of an infinite generator. """
    for i, x in enumerate(generator()):
        if i == n - 1:
            return x

As an example:

>>> nth_term(triangle_stream, 10)
55

But this nth_term function also works:

def nth_term(generator, n):
    """ Returns the nth term of an infinite generator. """
    x = generator()
    for i in range(n - 1):
        next(x)
    return next(x)

Which of these two nth_term functions is more efficient?

Two modifications that I would make right away would be:

• to use iterables instead of a generator function as parameter to allow for more genericity: this way I could use nth_term using either triangle_stream() or itertools.cycle('ABCD') as parameters;

• to index elements in base 0 instead of base 1, so nth_term(triangle_stream(), 0) would return 1 instead of looping indefinitely and nth_term(triangle_stream(), 1) would return 3 instead of 1.

Now, making the function accept any iterable, I would use something along the lines of the second version this will be closer to my final proposition that will raise StopIteration if there is not enough elements instead of silently returning None ("Errors should never pass silently."). But applying both suggestions would lead to:

def nth_term(iterable, n):
    """Returns the nth term of any iterable."""
    for i, x in enumerate(iterable):
        if i == n:
            return x

or

def nth_term(iterable, n):
    """Returns the nth term of any iterable."""
    x = iter(iterable)
    for i in range(n):
        next(x)
    return next(x)

Now, instead of manually calling next on the iterator in the second version, I would ask zip to do it for me. Chances are that the execution will spend less time in the Python interpreter and more in the C part of zip and so be faster:

def nth_term(iterable, n):
    """Returns the nth term of any iterable."""
    iterator = iter(iterable)
    for _ in zip(range(n), iterator):
        pass
    return next(iterator)

Lastly, I would use existing tools to perform the iteration for me: itertools.islice is exactly what we need. In fact, your very function is already listed in the itertools recipes:

def nth(iterable, n, default=None):
    "Returns the nth item or a default value"
    return next(islice(iterable, n, None), default)

I would however remove the default value, especially in the case of an infinite iterable:

def nth(iterable, n):
    "Returns the nth item or raise StopIteration"
    return next(islice(iterable, n, None))

A completely unrelated note, but I would move the yield x at the beginning of the loop in triangle_stream so that you don't have to write it twice:

def triangle_stream():
    """ Infinite stream of triangle numbers. """
    x = step = 1
    while True:
        yield x
        step += 1
        x += step

And since step starts at 1, and never stop incrementing one by one, I would use itertools.count(1):

def triangle_stream():
    """ Infinite stream of triangle numbers. """
    x = 0
    for step in itertools.count(1):
        x += step
        yield x

Now triangle_stream only accumulates values from itertools.count(1), well time to make use of itertools.accumulate:

def triangle_stream():
    """ Infinite stream of triangle numbers. """
    return itertools.accumulate(itertools.count(1))

or, to make it more obvious that you get a generator by calling triangle_stream:

def triangle_stream():
    """ Infinite stream of triangle numbers. """
    yield from itertools.accumulate(itertools.count(1))

StackExchange Code Review Q#145996, answer score: 9