class GeneratorContextManager(object): def __init__(self, gen): self.gen = gen def __context__(self): return self def __enter__(self): try: return self.gen.next() except StopIteration: raise RuntimeError("generator didn't yield") def __exit__(self, type, value, traceback): if type is None: try: self.gen.next() except StopIteration: return else: raise RuntimeError("generator didn't stop") else: try: self.gen.throw(type, value, traceback) except (type, StopIteration): return else: raise RuntimeError("generator caught exception") def contextmanager(func): def helper(*args, **kwds): return GeneratorContextManager(func(*args, **kwds)) return helper