Lib/contextlib.py

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