1 """Utilities for with-statement contexts. See PEP 343.""" 2 import abc 3 import sys 4 from collections import deque 5 from functools import wraps 6 7 __all__ = ["contextmanager", "closing", "AbstractContextManager", 8 "ContextDecorator", "ExitStack", "redirect_stdout", 9 "redirect_stderr", "suppress"] 10 11 12 class AbstractContextManager(abc.ABC): 13 14 """An abstract base class for context managers.""" 15 16 def __enter__(self): 17 """Return `self` upon entering the runtime context.""" 18 return self 19 20 @abc.abstractmethod 21 def __exit__(self, exc_type, exc_value, traceback): 22 """Raise any exception triggered within the runtime context.""" 23 return None 24 25 @classmethod 26 def __subclasshook__(cls, C): 27 if cls is AbstractContextManager: 28 if (any("__enter__" in B.__dict__ for B in C.__mro__) and 29 any("__exit__" in B.__dict__ for B in C.__mro__)): 30 return True 31 return NotImplemented 32 33 34 class ContextDecorator(object): 35 "A base class or mixin that enables context managers to work as decorators." 36 37 def _recreate_cm(self): 38 """Return a recreated instance of self. 39 40 Allows an otherwise one-shot context manager like 41 _GeneratorContextManager to support use as 42 a decorator via implicit recreation. 43 44 This is a private interface just for _GeneratorContextManager. 45 See issue #11647 for details. 46 """ 47 return self 48 49 def __call__(self, func): 50 @wraps(func) 51 def inner(*args, **kwds): 52 with self._recreate_cm(): 53 return func(*args, **kwds) 54 return inner 55 56 57 class _GeneratorContextManager(ContextDecorator, AbstractContextManager): 58 """Helper for @contextmanager decorator.""" 59 60 def __init__(self, func, args, kwds): 61 self.gen = func(*args, **kwds) 62 self.func, self.args, self.kwds = func, args, kwds 63 # Issue 19330: ensure context manager instances have good docstrings 64 doc = getattr(func, "__doc__", None) 65 if doc is None: 66 doc = type(self).__doc__ 67 self.__doc__ = doc 68 # Unfortunately, this still doesn't provide good help output when 69 # inspecting the created context manager instances, since pydoc 70 # currently bypasses the instance docstring and shows the docstring 71 # for the class instead. 72 # See http://bugs.python.org/issue19404 for more details. 73 74 def _recreate_cm(self): 75 # _GCM instances are one-shot context managers, so the 76 # CM must be recreated each time a decorated function is 77 # called 78 return self.__class__(self.func, self.args, self.kwds) 79 80 def __enter__(self): 81 try: 82 return next(self.gen) 83 except StopIteration: 84 raise RuntimeError("generator didn't yield") from None 85 86 def __exit__(self, type, value, traceback): 87 if type is None: 88 try: 89 next(self.gen) 90 except StopIteration: 91 return 92 else: 93 raise RuntimeError("generator didn't stop") 94 else: 95 if value is None: 96 # Need to force instantiation so we can reliably 97 # tell if we get the same exception back 98 value = type() 99 try: 100 self.gen.throw(type, value, traceback) 101 raise RuntimeError("generator didn't stop after throw()") 102 except StopIteration as exc: 103 # Suppress StopIteration *unless* it's the same exception that 104 # was passed to throw(). This prevents a StopIteration 105 # raised inside the "with" statement from being suppressed. 106 return exc is not value 107 except RuntimeError as exc: 108 # Don't re-raise the passed in exception. (issue27122) 109 if exc is value: 110 return False 111 # Likewise, avoid suppressing if a StopIteration exception 112 # was passed to throw() and later wrapped into a RuntimeError 113 # (see PEP 479). 114 if exc.__cause__ is value: 115 return False 116 raise 117 except: 118 # only re-raise if it's *not* the exception that was 119 # passed to throw(), because __exit__() must not raise 120 # an exception unless __exit__() itself failed. But throw() 121 # has to raise the exception to signal propagation, so this 122 # fixes the impedance mismatch between the throw() protocol 123 # and the __exit__() protocol. 124 # 125 if sys.exc_info()[1] is not value: 126 raise 127 128 129 def contextmanager(func): 130 """@contextmanager decorator. 131 132 Typical usage: 133 134 @contextmanager 135 def some_generator(<arguments>): 136 <setup> 137 try: 138 yield <value> 139 finally: 140 <cleanup> 141 142 This makes this: 143 144 with some_generator(<arguments>) as <variable>: 145 <body> 146 147 equivalent to this: 148 149 <setup> 150 try: 151 <variable> = <value> 152 <body> 153 finally: 154 <cleanup> 155 156 """ 157 @wraps(func) 158 def helper(*args, **kwds): 159 return _GeneratorContextManager(func, args, kwds) 160 return helper 161 162 163 class closing(AbstractContextManager): 164 """Context to automatically close something at the end of a block. 165 166 Code like this: 167 168 with closing(<module>.open(<arguments>)) as f: 169 <block> 170 171 is equivalent to this: 172 173 f = <module>.open(<arguments>) 174 try: 175 <block> 176 finally: 177 f.close() 178 179 """ 180 def __init__(self, thing): 181 self.thing = thing 182 def __enter__(self): 183 return self.thing 184 def __exit__(self, *exc_info): 185 self.thing.close() 186 187 188 class _RedirectStream(AbstractContextManager): 189 190 _stream = None 191 192 def __init__(self, new_target): 193 self._new_target = new_target 194 # We use a list of old targets to make this CM re-entrant 195 self._old_targets = [] 196 197 def __enter__(self): 198 self._old_targets.append(getattr(sys, self._stream)) 199 setattr(sys, self._stream, self._new_target) 200 return self._new_target 201 202 def __exit__(self, exctype, excinst, exctb): 203 setattr(sys, self._stream, self._old_targets.pop()) 204 205 206 class redirect_stdout(_RedirectStream): 207 """Context manager for temporarily redirecting stdout to another file. 208 209 # How to send help() to stderr 210 with redirect_stdout(sys.stderr): 211 help(dir) 212 213 # How to write help() to a file 214 with open('help.txt', 'w') as f: 215 with redirect_stdout(f): 216 help(pow) 217 """ 218 219 _stream = "stdout" 220 221 222 class redirect_stderr(_RedirectStream): 223 """Context manager for temporarily redirecting stderr to another file.""" 224 225 _stream = "stderr" 226 227 228 class suppress(AbstractContextManager): 229 """Context manager to suppress specified exceptions 230 231 After the exception is suppressed, execution proceeds with the next 232 statement following the with statement. 233 234 with suppress(FileNotFoundError): 235 os.remove(somefile) 236 # Execution still resumes here if the file was already removed 237 """ 238 239 def __init__(self, *exceptions): 240 self._exceptions = exceptions 241 242 def __enter__(self): 243 pass 244 245 def __exit__(self, exctype, excinst, exctb): 246 # Unlike isinstance and issubclass, CPython exception handling 247 # currently only looks at the concrete type hierarchy (ignoring 248 # the instance and subclass checking hooks). While Guido considers 249 # that a bug rather than a feature, it's a fairly hard one to fix 250 # due to various internal implementation details. suppress provides 251 # the simpler issubclass based semantics, rather than trying to 252 # exactly reproduce the limitations of the CPython interpreter. 253 # 254 # See http://bugs.python.org/issue12029 for more details 255 return exctype is not None and issubclass(exctype, self._exceptions) 256 257 258 # Inspired by discussions on http://bugs.python.org/issue13585 259 class ExitStack(AbstractContextManager): 260 """Context manager for dynamic management of a stack of exit callbacks 261 262 For example: 263 264 with ExitStack() as stack: 265 files = [stack.enter_context(open(fname)) for fname in filenames] 266 # All opened files will automatically be closed at the end of 267 # the with statement, even if attempts to open files later 268 # in the list raise an exception 269 270 """ 271 def __init__(self): 272 self._exit_callbacks = deque() 273 274 def pop_all(self): 275 """Preserve the context stack by transferring it to a new instance""" 276 new_stack = type(self)() 277 new_stack._exit_callbacks = self._exit_callbacks 278 self._exit_callbacks = deque() 279 return new_stack 280 281 def _push_cm_exit(self, cm, cm_exit): 282 """Helper to correctly register callbacks to __exit__ methods""" 283 def _exit_wrapper(*exc_details): 284 return cm_exit(cm, *exc_details) 285 _exit_wrapper.__self__ = cm 286 self.push(_exit_wrapper) 287 288 def push(self, exit): 289 """Registers a callback with the standard __exit__ method signature 290 291 Can suppress exceptions the same way __exit__ methods can. 292 293 Also accepts any object with an __exit__ method (registering a call 294 to the method instead of the object itself) 295 """ 296 # We use an unbound method rather than a bound method to follow 297 # the standard lookup behaviour for special methods 298 _cb_type = type(exit) 299 try: 300 exit_method = _cb_type.__exit__ 301 except AttributeError: 302 # Not a context manager, so assume its a callable 303 self._exit_callbacks.append(exit) 304 else: 305 self._push_cm_exit(exit, exit_method) 306 return exit # Allow use as a decorator 307 308 def callback(self, callback, *args, **kwds): 309 """Registers an arbitrary callback and arguments. 310 311 Cannot suppress exceptions. 312 """ 313 def _exit_wrapper(exc_type, exc, tb): 314 callback(*args, **kwds) 315 # We changed the signature, so using @wraps is not appropriate, but 316 # setting __wrapped__ may still help with introspection 317 _exit_wrapper.__wrapped__ = callback 318 self.push(_exit_wrapper) 319 return callback # Allow use as a decorator 320 321 def enter_context(self, cm): 322 """Enters the supplied context manager 323 324 If successful, also pushes its __exit__ method as a callback and 325 returns the result of the __enter__ method. 326 """ 327 # We look up the special methods on the type to match the with statement 328 _cm_type = type(cm) 329 _exit = _cm_type.__exit__ 330 result = _cm_type.__enter__(cm) 331 self._push_cm_exit(cm, _exit) 332 return result 333 334 def close(self): 335 """Immediately unwind the context stack""" 336 self.__exit__(None, None, None) 337 338 def __exit__(self, *exc_details): 339 received_exc = exc_details[0] is not None 340 341 # We manipulate the exception state so it behaves as though 342 # we were actually nesting multiple with statements 343 frame_exc = sys.exc_info()[1] 344 def _fix_exception_context(new_exc, old_exc): 345 # Context may not be correct, so find the end of the chain 346 while 1: 347 exc_context = new_exc.__context__ 348 if exc_context is old_exc: 349 # Context is already set correctly (see issue 20317) 350 return 351 if exc_context is None or exc_context is frame_exc: 352 break 353 new_exc = exc_context 354 # Change the end of the chain to point to the exception 355 # we expect it to reference 356 new_exc.__context__ = old_exc 357 358 # Callbacks are invoked in LIFO order to match the behaviour of 359 # nested context managers 360 suppressed_exc = False 361 pending_raise = False 362 while self._exit_callbacks: 363 cb = self._exit_callbacks.pop() 364 try: 365 if cb(*exc_details): 366 suppressed_exc = True 367 pending_raise = False 368 exc_details = (None, None, None) 369 except: 370 new_exc_details = sys.exc_info() 371 # simulate the stack of exceptions by setting the context 372 _fix_exception_context(new_exc_details[1], exc_details[1]) 373 pending_raise = True 374 exc_details = new_exc_details 375 if pending_raise: 376 try: 377 # bare "raise exc_details[1]" replaces our carefully 378 # set-up context 379 fixed_ctx = exc_details[1].__context__ 380 raise exc_details[1] 381 except BaseException: 382 exc_details[1].__context__ = fixed_ctx 383 raise 384 return received_exc and suppressed_exc 385
