From f73d2e70037ea0c76d03487a71790734344aa080 Mon Sep 17 00:00:00 2001 From: Serhiy Storchaka Date: Wed, 4 Feb 2026 13:20:18 +0200 Subject: [PATCH] gh-144386: Add support for descriptors in ExitStack and AsyncExitStack (#144420) __enter__(), __exit__(), __aenter__(), and __aexit__() can now be arbitrary descriptors, not only normal methods, for consistency with the "with" and "async with" statements. --- Doc/library/contextlib.rst | 13 +++ Doc/whatsnew/3.15.rst | 10 ++ Lib/contextlib.py | 99 ++++++++----------- Lib/test/test_contextlib.py | 69 +++++++++++++ Lib/test/test_contextlib_async.py | 72 ++++++++++++++ ...-02-03-14-16-49.gh-issue-144386.9Wa59r.rst | 4 + 6 files changed, 209 insertions(+), 58 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2026-02-03-14-16-49.gh-issue-144386.9Wa59r.rst diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst index d0fa645093a..f2e3c836cec 100644 --- a/Doc/library/contextlib.rst +++ b/Doc/library/contextlib.rst @@ -564,6 +564,10 @@ Functions and classes provided: Raises :exc:`TypeError` instead of :exc:`AttributeError` if *cm* is not a context manager. + .. versionchanged:: next + Added support for arbitrary descriptors :meth:`!__enter__` and + :meth:`!__exit__`. + .. method:: push(exit) Adds a context manager's :meth:`~object.__exit__` method to the callback stack. @@ -582,6 +586,9 @@ Functions and classes provided: The passed in object is returned from the function, allowing this method to be used as a function decorator. + .. versionchanged:: next + Added support for arbitrary descriptors :meth:`!__exit__`. + .. method:: callback(callback, /, *args, **kwds) Accepts an arbitrary callback function and arguments and adds it to @@ -639,11 +646,17 @@ Functions and classes provided: Raises :exc:`TypeError` instead of :exc:`AttributeError` if *cm* is not an asynchronous context manager. + .. versionchanged:: next + Added support for arbitrary descriptors :meth:`!__aenter__` and :meth:`!__aexit__`. + .. method:: push_async_exit(exit) Similar to :meth:`ExitStack.push` but expects either an asynchronous context manager or a coroutine function. + .. versionchanged:: next + Added support for arbitrary descriptors :meth:`!__aexit__`. + .. method:: push_async_callback(callback, /, *args, **kwds) Similar to :meth:`ExitStack.callback` but expects a coroutine function. diff --git a/Doc/whatsnew/3.15.rst b/Doc/whatsnew/3.15.rst index 7979933d7e8..2291228e721 100644 --- a/Doc/whatsnew/3.15.rst +++ b/Doc/whatsnew/3.15.rst @@ -548,6 +548,16 @@ concurrent.futures (Contributed by Jonathan Berg in :gh:`139486`.) +contextlib +---------- + +* Added support for arbitrary descriptors :meth:`!__enter__`, + :meth:`!__exit__`, :meth:`!__aenter__`, and :meth:`!__aexit__` in + :class:`~contextlib.ExitStack` and :class:`contextlib.AsyncExitStack`, for + consistency with the :keyword:`with` and :keyword:`async with` statements. + (Contributed by Serhiy Storchaka in :gh:`144386`.) + + dataclasses ----------- diff --git a/Lib/contextlib.py b/Lib/contextlib.py index 5b646fabca0..cac3e39eba8 100644 --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -5,7 +5,7 @@ import _collections_abc from collections import deque from functools import wraps -from types import MethodType, GenericAlias +from types import GenericAlias __all__ = ["asynccontextmanager", "contextmanager", "closing", "nullcontext", "AbstractContextManager", "AbstractAsyncContextManager", @@ -469,13 +469,23 @@ def __exit__(self, exctype, excinst, exctb): return False +def _lookup_special(obj, name, default): + # Follow the standard lookup behaviour for special methods. + from inspect import getattr_static, _descriptor_get + cls = type(obj) + try: + descr = getattr_static(cls, name) + except AttributeError: + return default + return _descriptor_get(descr, obj) + + +_sentinel = ['SENTINEL'] + + class _BaseExitStack: """A base class for ExitStack and AsyncExitStack.""" - @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): @@ -499,17 +509,8 @@ def push(self, exit): 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) + exit_method = _lookup_special(exit, '__exit__', exit) + self._push_exit_callback(exit_method) return exit # Allow use as a decorator. def enter_context(self, cm): @@ -518,17 +519,18 @@ def enter_context(self, cm): 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. - cls = type(cm) - try: - _enter = cls.__enter__ - _exit = cls.__exit__ - except AttributeError: + _enter = _lookup_special(cm, '__enter__', _sentinel) + if _enter is _sentinel: + cls = type(cm) 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) + f"not support the context manager protocol") + _exit = _lookup_special(cm, '__exit__', _sentinel) + if _exit is _sentinel: + cls = type(cm) + raise TypeError(f"'{cls.__module__}.{cls.__qualname__}' object does " + f"not support the context manager protocol") + result = _enter() + self._push_exit_callback(_exit) return result def callback(self, callback, /, *args, **kwds): @@ -544,11 +546,6 @@ def callback(self, callback, /, *args, **kwds): 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)) @@ -641,10 +638,6 @@ class AsyncExitStack(_BaseExitStack, AbstractAsyncContextManager): # connection later in the list raise an exception. """ - @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): @@ -657,16 +650,18 @@ async def enter_async_context(self, cm): If successful, also pushes its __aexit__ method as a callback and returns the result of the __aenter__ method. """ - cls = type(cm) - try: - _enter = cls.__aenter__ - _exit = cls.__aexit__ - except AttributeError: + _enter = _lookup_special(cm, '__aenter__', _sentinel) + if _enter is _sentinel: + cls = type(cm) 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) + f"not support the asynchronous context manager protocol") + _exit = _lookup_special(cm, '__aexit__', _sentinel) + if _exit is _sentinel: + cls = type(cm) + raise TypeError(f"'{cls.__module__}.{cls.__qualname__}' object does " + f"not support the asynchronous context manager protocol") + result = await _enter() + self._push_exit_callback(_exit, False) return result def push_async_exit(self, exit): @@ -677,14 +672,8 @@ def push_async_exit(self, exit): 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) + exit_method = _lookup_special(exit, '__aexit__', exit) + self._push_exit_callback(exit_method, False) return exit # Allow use as a decorator def push_async_callback(self, callback, /, *args, **kwds): @@ -704,12 +693,6 @@ 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 diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py index 6a3329fa5aa..1fd8b3cb18c 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -788,6 +788,75 @@ def _exit(): result.append(2) self.assertEqual(result, [1, 2, 3, 4]) + def test_enter_context_classmethod(self): + class TestCM: + @classmethod + def __enter__(cls): + result.append(('enter', cls)) + @classmethod + def __exit__(cls, *exc_details): + result.append(('exit', cls, *exc_details)) + + cm = TestCM() + result = [] + with self.exit_stack() as stack: + stack.enter_context(cm) + self.assertEqual(result, [('enter', TestCM)]) + self.assertEqual(result, [('enter', TestCM), + ('exit', TestCM, None, None, None)]) + + result = [] + with self.exit_stack() as stack: + stack.push(cm) + self.assertEqual(result, []) + self.assertEqual(result, [('exit', TestCM, None, None, None)]) + + def test_enter_context_staticmethod(self): + class TestCM: + @staticmethod + def __enter__(): + result.append('enter') + @staticmethod + def __exit__(*exc_details): + result.append(('exit', *exc_details)) + + cm = TestCM() + result = [] + with self.exit_stack() as stack: + stack.enter_context(cm) + self.assertEqual(result, ['enter']) + self.assertEqual(result, ['enter', ('exit', None, None, None)]) + + result = [] + with self.exit_stack() as stack: + stack.push(cm) + self.assertEqual(result, []) + self.assertEqual(result, [('exit', None, None, None)]) + + def test_enter_context_slots(self): + class TestCM: + __slots__ = ('__enter__', '__exit__') + def __init__(self): + def enter(): + result.append('enter') + def exit(*exc_details): + result.append(('exit', *exc_details)) + self.__enter__ = enter + self.__exit__ = exit + + cm = TestCM() + result = [] + with self.exit_stack() as stack: + stack.enter_context(cm) + self.assertEqual(result, ['enter']) + self.assertEqual(result, ['enter', ('exit', None, None, None)]) + + result = [] + with self.exit_stack() as stack: + stack.push(cm) + self.assertEqual(result, []) + self.assertEqual(result, [('exit', None, None, None)]) + def test_enter_context_errors(self): class LacksEnterAndExit: pass diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py index dcd00720379..248d32d6152 100644 --- a/Lib/test/test_contextlib_async.py +++ b/Lib/test/test_contextlib_async.py @@ -641,6 +641,78 @@ async def _exit(): self.assertEqual(result, [1, 2, 3, 4]) + @_async_test + async def test_enter_async_context_classmethod(self): + class TestCM: + @classmethod + async def __aenter__(cls): + result.append(('enter', cls)) + @classmethod + async def __aexit__(cls, *exc_details): + result.append(('exit', cls, *exc_details)) + + cm = TestCM() + result = [] + async with self.exit_stack() as stack: + await stack.enter_async_context(cm) + self.assertEqual(result, [('enter', TestCM)]) + self.assertEqual(result, [('enter', TestCM), + ('exit', TestCM, None, None, None)]) + + result = [] + async with self.exit_stack() as stack: + stack.push_async_exit(cm) + self.assertEqual(result, []) + self.assertEqual(result, [('exit', TestCM, None, None, None)]) + + @_async_test + async def test_enter_async_context_staticmethod(self): + class TestCM: + @staticmethod + async def __aenter__(): + result.append('enter') + @staticmethod + async def __aexit__(*exc_details): + result.append(('exit', *exc_details)) + + cm = TestCM() + result = [] + async with self.exit_stack() as stack: + await stack.enter_async_context(cm) + self.assertEqual(result, ['enter']) + self.assertEqual(result, ['enter', ('exit', None, None, None)]) + + result = [] + async with self.exit_stack() as stack: + stack.push_async_exit(cm) + self.assertEqual(result, []) + self.assertEqual(result, [('exit', None, None, None)]) + + @_async_test + async def test_enter_async_context_slots(self): + class TestCM: + __slots__ = ('__aenter__', '__aexit__') + def __init__(self): + async def enter(): + result.append('enter') + async def exit(*exc_details): + result.append(('exit', *exc_details)) + self.__aenter__ = enter + self.__aexit__ = exit + + cm = TestCM() + result = [] + async with self.exit_stack() as stack: + await stack.enter_async_context(cm) + self.assertEqual(result, ['enter']) + self.assertEqual(result, ['enter', ('exit', None, None, None)]) + + result = [] + async with self.exit_stack() as stack: + stack.push_async_exit(cm) + self.assertEqual(result, []) + self.assertEqual(result, [('exit', None, None, None)]) + @_async_test async def test_enter_async_context_errors(self): class LacksEnterAndExit: diff --git a/Misc/NEWS.d/next/Library/2026-02-03-14-16-49.gh-issue-144386.9Wa59r.rst b/Misc/NEWS.d/next/Library/2026-02-03-14-16-49.gh-issue-144386.9Wa59r.rst new file mode 100644 index 00000000000..6e60eeba208 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-02-03-14-16-49.gh-issue-144386.9Wa59r.rst @@ -0,0 +1,4 @@ +Add support for arbitrary descriptors :meth:`!__enter__`, :meth:`!__exit__`, +:meth:`!__aenter__`, and :meth:`!__aexit__` in :class:`contextlib.ExitStack` +and :class:`contextlib.AsyncExitStack`, for consistency with the +:keyword:`with` and :keyword:`async with` statements.