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