Lib/test/contextmanager.py

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