From 0e9194f8ce9b7d070c992ff5d4985cac4d4fe250 Mon Sep 17 00:00:00 2001 From: Jack O'Connor Date: Wed, 3 Oct 2018 15:27:45 -0400 Subject: [PATCH] add the AsyncExitStack backport, to restore support for Python 3.5 and 3.6 --- peru/async_exit_stack.py | 218 +++++++++++++++++++++++++++++++++++++++ peru/plugin.py | 5 +- 2 files changed, 221 insertions(+), 2 deletions(-) create mode 100644 peru/async_exit_stack.py diff --git a/peru/async_exit_stack.py b/peru/async_exit_stack.py new file mode 100644 index 0000000..da4d311 --- /dev/null +++ b/peru/async_exit_stack.py @@ -0,0 +1,218 @@ +# This definition of AsyncExitStack is copied entirely from +# https://github.com/sorcio/async_exit_stack, in order to support Python 3.5 +# and 3.6. +# +# flake8: noqa + +from collections import deque +import sys +from types import MethodType + + +class AsyncExitStack: + """Async context manager for dynamic management of a stack of exit + callbacks. + For example: + async with AsyncExitStack() as stack: + connections = [await stack.enter_async_context(get_connection()) + for i in range(5)] + # All opened connections will automatically be released at the + # end of the async with statement, even if attempts to open a + # connection later in the list raise an exception. + """ + + ### _BaseExitStack staticmethods + + @staticmethod + def _create_exit_wrapper(cm, cm_exit): + return MethodType(cm_exit, cm) + + @staticmethod + def _create_cb_wrapper(callback, *args, **kwds): + def _exit_wrapper(exc_type, exc, tb): + callback(*args, **kwds) + return _exit_wrapper + + ### AsyncExitStack staticmethods + + @staticmethod + def _create_async_exit_wrapper(cm, cm_exit): + return MethodType(cm_exit, cm) + + @staticmethod + def _create_async_cb_wrapper(callback, *args, **kwds): + async def _exit_wrapper(exc_type, exc, tb): + await callback(*args, **kwds) + return _exit_wrapper + + ### _BaseExitStack methods + + def __init__(self): + self._exit_callbacks = deque() + + def pop_all(self): + """Preserve the context stack by transferring it to a new instance.""" + new_stack = type(self)() + new_stack._exit_callbacks = self._exit_callbacks + self._exit_callbacks = deque() + return new_stack + + def push(self, exit): + """Registers a callback with the standard __exit__ method signature. + Can suppress exceptions the same way __exit__ method can. + Also accepts any object with an __exit__ method (registering a call + to the method instead of the object itself). + """ + # We use an unbound method rather than a bound method to follow + # the standard lookup behaviour for special methods. + _cb_type = type(exit) + + try: + exit_method = _cb_type.__exit__ + except AttributeError: + # Not a context manager, so assume it's a callable. + self._push_exit_callback(exit) + else: + self._push_cm_exit(exit, exit_method) + return exit # Allow use as a decorator. + + def enter_context(self, cm): + """Enters the supplied context manager. + If successful, also pushes its __exit__ method as a callback and + returns the result of the __enter__ method. + """ + # We look up the special methods on the type to match the with + # statement. + _cm_type = type(cm) + _exit = _cm_type.__exit__ + result = _cm_type.__enter__(cm) + self._push_cm_exit(cm, _exit) + return result + + def callback(self, callback, *args, **kwds): + """Registers an arbitrary callback and arguments. + Cannot suppress exceptions. + """ + _exit_wrapper = self._create_cb_wrapper(callback, *args, **kwds) + + # We changed the signature, so using @wraps is not appropriate, but + # setting __wrapped__ may still help with introspection. + _exit_wrapper.__wrapped__ = callback + self._push_exit_callback(_exit_wrapper) + return callback # Allow use as a decorator + + def _push_cm_exit(self, cm, cm_exit): + """Helper to correctly register callbacks to __exit__ methods.""" + _exit_wrapper = self._create_exit_wrapper(cm, cm_exit) + self._push_exit_callback(_exit_wrapper, True) + + def _push_exit_callback(self, callback, is_sync=True): + self._exit_callbacks.append((is_sync, callback)) + + ### AsyncExitStack methods + + async def enter_async_context(self, cm): + """Enters the supplied async context manager. + If successful, also pushes its __aexit__ method as a callback and + returns the result of the __aenter__ method. + """ + _cm_type = type(cm) + _exit = _cm_type.__aexit__ + result = await _cm_type.__aenter__(cm) + self._push_async_cm_exit(cm, _exit) + return result + + def push_async_exit(self, exit): + """Registers a coroutine function with the standard __aexit__ method + signature. + Can suppress exceptions the same way __aexit__ method can. + Also accepts any object with an __aexit__ method (registering a call + to the method instead of the object itself). + """ + _cb_type = type(exit) + try: + exit_method = _cb_type.__aexit__ + except AttributeError: + # Not an async context manager, so assume it's a coroutine function + self._push_exit_callback(exit, False) + else: + self._push_async_cm_exit(exit, exit_method) + return exit # Allow use as a decorator + + def push_async_callback(self, callback, *args, **kwds): + """Registers an arbitrary coroutine function and arguments. + Cannot suppress exceptions. + """ + _exit_wrapper = self._create_async_cb_wrapper(callback, *args, **kwds) + + # We changed the signature, so using @wraps is not appropriate, but + # setting __wrapped__ may still help with introspection. + _exit_wrapper.__wrapped__ = callback + self._push_exit_callback(_exit_wrapper, False) + return callback # Allow use as a decorator + + async def aclose(self): + """Immediately unwind the context stack.""" + await self.__aexit__(None, None, None) + + def _push_async_cm_exit(self, cm, cm_exit): + """Helper to correctly register coroutine function to __aexit__ + method.""" + _exit_wrapper = self._create_async_exit_wrapper(cm, cm_exit) + self._push_exit_callback(_exit_wrapper, False) + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc_details): + received_exc = exc_details[0] is not None + + # We manipulate the exception state so it behaves as though + # we were actually nesting multiple with statements + frame_exc = sys.exc_info()[1] + def _fix_exception_context(new_exc, old_exc): + # Context may not be correct, so find the end of the chain + while 1: + exc_context = new_exc.__context__ + if exc_context is old_exc: + # Context is already set correctly (see issue 20317) + return + if exc_context is None or exc_context is frame_exc: + break + new_exc = exc_context + # Change the end of the chain to point to the exception + # we expect it to reference + new_exc.__context__ = old_exc + + # Callbacks are invoked in LIFO order to match the behaviour of + # nested context managers + suppressed_exc = False + pending_raise = False + while self._exit_callbacks: + is_sync, cb = self._exit_callbacks.pop() + try: + if is_sync: + cb_suppress = cb(*exc_details) + else: + cb_suppress = await cb(*exc_details) + + if cb_suppress: + suppressed_exc = True + pending_raise = False + exc_details = (None, None, None) + except: + new_exc_details = sys.exc_info() + # simulate the stack of exceptions by setting the context + _fix_exception_context(new_exc_details[1], exc_details[1]) + pending_raise = True + exc_details = new_exc_details + if pending_raise: + try: + # bare "raise exc_details[1]" replaces our carefully + # set-up context + fixed_ctx = exc_details[1].__context__ + raise exc_details[1] + except BaseException: + exc_details[1].__context__ = fixed_ctx + raise + return received_exc and suppressed_exc diff --git a/peru/plugin.py b/peru/plugin.py index 43a5351..a7b7bba 100644 --- a/peru/plugin.py +++ b/peru/plugin.py @@ -7,6 +7,7 @@ import tempfile import yaml from .async_helpers import create_subprocess_with_handle +from .async_exit_stack import AsyncExitStack from . import cache from . import compat from .compat import makedirs @@ -60,7 +61,7 @@ async def _plugin_job(plugin_context, module_type, module_fields, command, env, display_handle): # We take several locks and other context managers in here. Using an # AsyncExitStack saves us from indentation hell. - async with contextlib.AsyncExitStack() as stack: + async with AsyncExitStack() as stack: definition = _get_plugin_definition(module_type, module_fields, command) exe = _get_plugin_exe(definition, command) @@ -188,7 +189,7 @@ def _plugin_env(plugin_context, plugin_definition, module_fields, command, def _noop_lock(): - return contextlib.AsyncExitStack() # a no-op context manager + return AsyncExitStack() # a no-op context manager def _plugin_cache_lock(plugin_context, definition, module_fields):