14adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao"""Utilities for with-statement contexts.  See PEP 343."""
24adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
34adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport sys
44adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaofrom functools import wraps
54adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaofrom warnings import warn
64adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
74adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao__all__ = ["contextmanager", "nested", "closing"]
84adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
94adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoclass GeneratorContextManager(object):
104adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    """Helper for @contextmanager decorator."""
114adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
124adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def __init__(self, gen):
134adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.gen = gen
144adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
154adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def __enter__(self):
164adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        try:
174adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            return self.gen.next()
184adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        except StopIteration:
194adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            raise RuntimeError("generator didn't yield")
204adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
214adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def __exit__(self, type, value, traceback):
224adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        if type is None:
234adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            try:
244adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                self.gen.next()
254adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            except StopIteration:
264adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                return
274adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            else:
284adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                raise RuntimeError("generator didn't stop")
294adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        else:
304adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            if value is None:
314adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # Need to force instantiation so we can reliably
324adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # tell if we get the same exception back
334adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                value = type()
344adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            try:
354adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                self.gen.throw(type, value, traceback)
364adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                raise RuntimeError("generator didn't stop after throw()")
374adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            except StopIteration, exc:
384adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # Suppress the exception *unless* it's the same exception that
394adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # was passed to throw().  This prevents a StopIteration
404adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # raised inside the "with" statement from being suppressed
414adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                return exc is not value
424adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            except:
434adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # only re-raise if it's *not* the exception that was
444adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # passed to throw(), because __exit__() must not raise
454adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # an exception unless __exit__() itself failed.  But throw()
464adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # has to raise the exception to signal propagation, so this
474adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # fixes the impedance mismatch between the throw() protocol
484adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                # and the __exit__() protocol.
494adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                #
504adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                if sys.exc_info()[1] is not value:
514adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                    raise
524adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
534adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
544adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaodef contextmanager(func):
554adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    """@contextmanager decorator.
564adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
574adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    Typical usage:
584adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
594adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        @contextmanager
604adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        def some_generator(<arguments>):
614adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            <setup>
624adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            try:
634adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                yield <value>
644adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            finally:
654adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                <cleanup>
664adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
674adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    This makes this:
684adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
694adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        with some_generator(<arguments>) as <variable>:
704adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            <body>
714adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
724adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    equivalent to this:
734adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
744adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        <setup>
754adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        try:
764adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            <variable> = <value>
774adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            <body>
784adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        finally:
794adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            <cleanup>
804adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
814adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    """
824adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    @wraps(func)
834adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def helper(*args, **kwds):
844adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        return GeneratorContextManager(func(*args, **kwds))
854adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    return helper
864adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
874adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
884adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao@contextmanager
894adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaodef nested(*managers):
904adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    """Combine multiple context managers into a single nested context manager.
914adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
924adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao   This function has been deprecated in favour of the multiple manager form
934adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao   of the with statement.
944adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
954adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao   The one advantage of this function over the multiple manager form of the
964adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao   with statement is that argument unpacking allows it to be
974adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao   used with a variable number of context managers as follows:
984adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
994adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao      with nested(*managers):
1004adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao          do_something()
1014adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1024adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    """
1034adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    warn("With-statements now directly support multiple context managers",
1044adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao         DeprecationWarning, 3)
1054adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    exits = []
1064adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    vars = []
1074adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    exc = (None, None, None)
1084adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    try:
1094adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        for mgr in managers:
1104adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            exit = mgr.__exit__
1114adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            enter = mgr.__enter__
1124adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            vars.append(enter())
1134adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            exits.append(exit)
1144adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        yield vars
1154adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    except:
1164adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        exc = sys.exc_info()
1174adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    finally:
1184adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        while exits:
1194adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            exit = exits.pop()
1204adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            try:
1214adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                if exit(*exc):
1224adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                    exc = (None, None, None)
1234adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            except:
1244adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao                exc = sys.exc_info()
1254adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        if exc != (None, None, None):
1264adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            # Don't rely on sys.exc_info() still containing
1274adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            # the right information. Another exception may
1284adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            # have been raised and caught by an exit method
1294adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            raise exc[0], exc[1], exc[2]
1304adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1314adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1324adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoclass closing(object):
1334adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    """Context to automatically close something at the end of a block.
1344adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1354adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    Code like this:
1364adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1374adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        with closing(<module>.open(<arguments>)) as f:
1384adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            <block>
1394adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1404adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    is equivalent to this:
1414adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1424adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f = <module>.open(<arguments>)
1434adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        try:
1444adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            <block>
1454adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        finally:
1464adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            f.close()
1474adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1484adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    """
1494adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def __init__(self, thing):
1504adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.thing = thing
1514adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def __enter__(self):
1524adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        return self.thing
1534adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def __exit__(self, *exc_info):
1544adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.thing.close()
155