diff options
Diffstat (limited to 'Lib/contextlib.py')
| -rw-r--r-- | Lib/contextlib.py | 23 |
1 files changed, 17 insertions, 6 deletions
diff --git a/Lib/contextlib.py b/Lib/contextlib.py index 1a8ef61..004d103 100644 --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -473,9 +473,14 @@ class _BaseExitStack: """ # 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) + cls = type(cm) + try: + _enter = cls.__enter__ + _exit = cls.__exit__ + except AttributeError: + raise TypeError(f"'{cls.__module__}.{cls.__qualname__}' object does " + f"not support the context manager protocol") from None + result = _enter(cm) self._push_cm_exit(cm, _exit) return result @@ -600,9 +605,15 @@ class AsyncExitStack(_BaseExitStack, AbstractAsyncContextManager): 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) + cls = type(cm) + try: + _enter = cls.__aenter__ + _exit = cls.__aexit__ + except AttributeError: + raise TypeError(f"'{cls.__module__}.{cls.__qualname__}' object does " + f"not support the asynchronous context manager protocol" + ) from None + result = await _enter(cm) self._push_async_cm_exit(cm, _exit) return result |