Lib/contextlib.py

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