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() 155