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