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             try:
  29                 self.gen.throw(type, value, traceback)
  30                 raise RuntimeError("generator didn't stop after throw()")
  31             except StopIteration, exc:
  32                 # Suppress the exception *unless* it's the same exception that
  33                 # was passed to throw().  This prevents a StopIteration
  34                 # raised inside the "with" statement from being suppressed
  35                 return exc is not value
  36             except:
  37                 # only re-raise if it's *not* the exception that was
  38                 # passed to throw(), because __exit__() must not raise
  39                 # an exception unless __exit__() itself failed.  But throw()
  40                 # has to raise the exception to signal propagation, so this
  41                 # fixes the impedance mismatch between the throw() protocol
  42                 # and the __exit__() protocol.
  43                 #
  44                 if sys.exc_info()[1] is not value:
  45                     raise
  46
  47
  48 def contextmanager(func):
  49     """@contextmanager decorator.
  50
  51     Typical usage:
  52
  53         @contextmanager
  54         def some_generator(<arguments>):
  55             <setup>
  56             try:
  57                 yield <value>
  58             finally:
  59                 <cleanup>
  60
  61     This makes this:
  62
  63         with some_generator(<arguments>) as <variable>:
  64             <body>
  65
  66     equivalent to this:
  67
  68         <setup>
  69         try:
  70             <variable> = <value>
  71             <body>
  72         finally:
  73             <cleanup>
  74
  75     """
  76     def helper(*args, **kwds):
  77         return GeneratorContextManager(func(*args, **kwds))
  78     try:
  79         helper.__name__ = func.__name__
  80         helper.__doc__ = func.__doc__
  81         helper.__dict__ = func.__dict__
  82     except:
  83         pass
  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         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()