Lib/contextlib.py

   1 """Utilities for with-statement contexts.  See PEP 343."""
   2
   3 import sys
   4 from functools import wraps
   5 from warnings import warn
   6
   7 __all__ = ["contextmanager", "nested", "closing"]
   8
   9 class GeneratorContextManager(object):
  10     """Helper for @contextmanager decorator."""
  11
  12     def __init__(self, gen):
  13         self.gen = gen
  14
  15     def __enter__(self):
  16         try:
  17             return next(self.gen)
  18         except StopIteration:
  19             raise RuntimeError("generator didn't yield")
  20
  21     def __exit__(self, type, value, traceback):
  22         if type is None:
  23             try:
  24                 next(self.gen)
  25             except StopIteration:
  26                 return
  27             else:
  28                 raise RuntimeError("generator didn't stop")
  29         else:
  30             if value is None:
  31                 # Need to force instantiation so we can reliably
  32                 # tell if we get the same exception back
  33                 value = type()
  34             try:
  35                 self.gen.throw(type, value, traceback)
  36                 raise RuntimeError("generator didn't stop after throw()")
  37             except StopIteration as exc:
  38                 # Suppress the exception *unless* it's the same exception that
  39                 # was passed to throw().  This prevents a StopIteration
  40                 # raised inside the "with" statement from being suppressed
  41                 return exc is not value
  42             except:
  43                 # only re-raise if it's *not* the exception that was
  44                 # passed to throw(), because __exit__() must not raise
  45                 # an exception unless __exit__() itself failed.  But throw()
  46                 # has to raise the exception to signal propagation, so this
  47                 # fixes the impedance mismatch between the throw() protocol
  48                 # and the __exit__() protocol.
  49                 #
  50                 if sys.exc_info()[1] is not value:
  51                     raise
  52
  53
  54 def contextmanager(func):
  55     """@contextmanager decorator.
  56
  57     Typical usage:
  58
  59         @contextmanager
  60         def some_generator(<arguments>):
  61             <setup>
  62             try:
  63                 yield <value>
  64             finally:
  65                 <cleanup>
  66
  67     This makes this:
  68
  69         with some_generator(<arguments>) as <variable>:
  70             <body>
  71
  72     equivalent to this:
  73
  74         <setup>
  75         try:
  76             <variable> = <value>
  77             <body>
  78         finally:
  79             <cleanup>
  80
  81     """
  82     @wraps(func)
  83     def helper(*args, **kwds):
  84         return GeneratorContextManager(func(*args, **kwds))
  85     return helper
  86
  87
  88 @contextmanager
  89 def nested(*managers):
  90     """Combine multiple context managers into a single nested context manager.
  91
  92    This function has been deprecated in favour of the multiple manager form
  93    of the with statement.
  94
  95    The one advantage of this function over the multiple manager form of the
  96    with statement is that argument unpacking allows it to be
  97    used with a variable number of context managers as follows:
  98
  99       with nested(*managers):
 100           do_something()
 101
 102     """
 103     warn("With-statements now directly support multiple context managers",
 104         DeprecationWarning, 3)
 105     exits = []
 106     vars = []
 107     exc = (None, None, None)
 108     try:
 109         for mgr in managers:
 110             exit = mgr.__exit__
 111             enter = mgr.__enter__
 112             vars.append(enter())
 113             exits.append(exit)
 114         yield vars
 115     except:
 116         exc = sys.exc_info()
 117     finally:
 118         while exits:
 119             exit = exits.pop()
 120             try:
 121                 if exit(*exc):
 122                     exc = (None, None, None)
 123             except:
 124                 exc = sys.exc_info()
 125         if exc != (None, None, None):
 126             # Don't rely on sys.exc_info() still containing
 127             # the right information. Another exception may
 128             # have been raised and caught by an exit method
 129             # exc[1] already has the __traceback__ attribute populated
 130             raise exc[1]
 131
 132
 133 class closing(object):
 134     """Context to automatically close something at the end of a block.
 135
 136     Code like this:
 137
 138         with closing(<module>.open(<arguments>)) as f:
 139             <block>
 140
 141     is equivalent to this:
 142
 143         f = <module>.open(<arguments>)
 144         try:
 145             <block>
 146         finally:
 147             f.close()
 148
 149     """
 150     def __init__(self, thing):
 151         self.thing = thing
 152     def __enter__(self):
 153         return self.thing
 154     def __exit__(self, *exc_info):
 155         self.thing.close()