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