将 python 生成器改造为上下文管理器

class Test:
    def __enter__(self):
        print('now in __enter__')
        return 'Hello World'

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('now exit')
        return True

if __name__ == '__main__':
    test = Test()
    with test as teststr:
        print(teststr)
    print('end of main')

import contextlib

class Test:
    @contextlib.contextmanager
    def contextmanager(self):
        print('now in __enter__')
        yield 'Hello World'
        print('now exit')
        return True

if __name__ == '__main__':
    test = Test()
    with test.contextmanager() as teststr:
        print(teststr)
    print('end of main')

def contextmanager(func):
    @wraps(func)
    def helper(*args, **kwds):
        return _GeneratorContextManager(func, args, kwds)
    return helper

class _GeneratorContextManager():
    def __init__(self, func, args, kwds):
        self.gen = func(*args, **kwds)
        self.func, self.args, self.kwds = func, args, kwds
        # Issue 19330: ensure context manager instances have good docstrings
        doc = getattr(func, "__doc__", None)
        if doc is None:
            doc = type(self).__doc__
        self.__doc__ = doc

    def __enter__(self):
        try:
            return next(self.gen)
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

    def __exit__(self, type, value, traceback):
        if type is None:
            try:
                next(self.gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("generator didn't stop")
        else:
            if value is None:
                # Need to force instantiation so we can reliably
                # tell if we get the same exception back
                value = type()
            try:
                self.gen.throw(type, value, traceback)
            except StopIteration as exc:
                # Suppress StopIteration *unless* it's the same exception that
                # was passed to throw().  This prevents a StopIteration
                # raised inside the "with" statement from being suppressed.
                return exc is not value
            except RuntimeError as exc:
                # Don't re-raise the passed in exception. (issue27122)
                if exc is value:
                    return False
                # Likewise, avoid suppressing if a StopIteration exception
                # was passed to throw() and later wrapped into a RuntimeError
                # (see PEP 479).
                if type is StopIteration and exc.__cause__ is value:
                    return False
                raise
            except:
                # only re-raise if it's *not* the exception that was
                # passed to throw(), because __exit__() must not raise
                # an exception unless __exit__() itself failed.  But throw()
                # has to raise the exception to signal propagation, so this
                # fixes the impedance mismatch between the throw() protocol
                # and the __exit__() protocol.
                #
                if sys.exc_info()[1] is value:
                    return False
                raise
            raise RuntimeError("generator didn't stop after throw()")

import contextlib

class Test:
    @contextlib.contextmanager
    def contextmanager(self):
        print('now in __enter__')
        yield self.raiseexc(1)
        print('now exit')
        return True

    def raiseexc(self, param):
        if param < 5:
            raise Exception('test exception')

if __name__ == '__main__':
    test = Test()
    with test.contextmanager() as teststr:
        print(teststr)
    print('end of main')

import contextlib

class Test:
    @contextlib.contextmanager
    def contextmanager(self):
        print('now in __enter__')
        try:
            yield self.raiseexc(1)
        except Exception:
            print('exception happened')
        print('now exit')
        return True

    def raiseexc(self, param):
        if param < 5:
            raise Exception('test exception')

if __name__ == '__main__':
    test = Test()
    with test.contextmanager() as teststr:
        print(teststr)
    print('end of main')