# -*- coding: utf-8 -*-

"""Tail call optimization for generators."""

__all__ = ["gtco", "gtrampolined"]

from functools import wraps

from inspect import isgenerator

def gtco(generator):

"""Low-level function: run a generator with TCO enabled.

In the generator, use ``return`` to tail-chain to the next generator.

Example::

def march():

yield 1

yield 2

return march() # tail-chain to a new instance of itself

assert tuple(take(6, gtco(march()))) == (1, 2, 1, 2, 1, 2)

last(take(10000, gtco(march()))) # no crash

"""

while True: # trampoline

x = yield from generator # yield stuff, get final result (return ...)

# don't let the TCO jump target bring along its trampoline if it has one

if isinstance(x, _TrampolinedGenerator):

x = x.g

if isgenerator(x):

generator = x

else:

# usually the return value is None, but allow for an iterable

try:

yield from x # the last batch!

except TypeError:

return x # passthrough

def gtrampolined(gfunc):

"""Decorator for generator functions (i.e. definitions of generators).

Decorating the definition avoids the need to use ``gtco`` at call time.

Example::

@gtrampolined

def ones():

yield 1

return ones()

assert tuple(take(10, ones())) == (1,) * 10

last(take(10000, ones())) # no crash

"""

@wraps(gfunc)

def trampolining_gfunc(*args, **kwargs):

generator = gfunc(*args, **kwargs)

return _TrampolinedGenerator(generator) # inject a trampoline

return trampolining_gfunc

class _TrampolinedGenerator:

"""Wrapper to inject the gtco() call to the generator g returned by gfunc."""

def __init__(self, g):

self.g = g

def __iter__(self):

return gtco(iter(self.g)) # start the trampoline

# no __next__, because __iter__ redirects;

# this wrapper is never actually iterated over.