HERE = object()

class GOTO(object):
    def __init__(self, gen, val=None):
        self.data = gen, val

class GOSUB(GOTO): pass

class RETURN(StopIteration):
    def __init__(self, val=None):
        self.val = val


def co(cogen, val=None):
    def gen(*args, **kwargs):
        gen = cogen(*args, **kwargs)
        val = None
        callstack = []
        while True:
            try:
                ret = gen.send(val) if val else gen.next()
            except StopIteration, e:
                if callstack:
                    gen, val = callstack.pop(), getattr(e, 'val', None)
                    continue
                raise
            if isinstance(ret, GOTO):
                if isinstance(ret, GOSUB):
                    callstack.append(gen)
                gen, val = ret.data
            elif ret is HERE:
                val = gen
            else:
                val = (yield ret)
    return gen

def trace(gen):
    for x in enumerate(gen):
        print "%s\t%s" % x


# Example of 'co-generators', which pass control to each other

def cofoo():
    self = (yield HERE)
    other = cobar(self)
    yield 'foo start'
    msg_from_cobar = yield GOTO(other)
    yield 'foo middle, cobar says "%s"' % msg_from_cobar
    yield GOTO(other)
    yield 'foo end'

def cobar(other):
    yield 'bar start'
    yield GOTO(other, 'hello')
    yield 'bar middle'
    yield GOTO(other)
    # We'll never get past here...
    yield 'bar end'
    
# Example of nesting of cogenerators

@co
def main_gen():
   yield 'start main'
   result = yield GOSUB(nested_cogen())
   yield 'returned %s' % result
   yield GOTO(nested_cogen())
   yield 'end main' # Never happens

def nested_cogen():
   yield 'nested'
   raise RETURN(1)


# Example of recursive generator

def coflatten(x):
   if isinstance(x, (list, tuple)):
      for i in x:
         yield GOSUB(coflatten(i))
   else:
      yield x

flatten = co(coflatten)