1"""Utilities for with-statement contexts.  See PEP 343."""
2import abc
3import sys
4from collections import deque
5from functools import wraps
6
7__all__ = ["contextmanager", "closing", "AbstractContextManager",
8           "ContextDecorator", "ExitStack", "redirect_stdout",
9           "redirect_stderr", "suppress"]
10
11
12class 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
34class 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
57class _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
129def 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
163class 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
188class _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
206class 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
222class redirect_stderr(_RedirectStream):
223    """Context manager for temporarily redirecting stderr to another file."""
224
225    _stream = "stderr"
226
227
228class 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
259class 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