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 self.gen.next()
  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                 self.gen.next()
  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, 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             raise exc[0], exc[1], exc[2]
 130
 131
 132 class closing(object):
 133     """Context to automatically close something at the end of a block.
 134
 135     Code like this:
 136
 137         with closing(<module>.open(<arguments>)) as f:
 138             <block>
 139
 140     is equivalent to this:
 141
 142         f = <module>.open(<arguments>)
 143         try:
 144             <block>
 145         finally:
 146             f.close()
 147
 148     """
 149     def __init__(self, thing):
 150         self.thing = thing
 151     def __enter__(self):
 152         return self.thing
 153     def __exit__(self, *exc_info):
 154         self.thing.close()