@@ -5,7 +5,7 @@
55import_collections_abc
66fromcollectionsimportdeque
77fromfunctoolsimportwraps
8-fromtypesimportMethodType, GenericAlias
8+fromtypesimportGenericAlias
991010__all__= ["asynccontextmanager", "contextmanager", "closing", "nullcontext",
1111"AbstractContextManager", "AbstractAsyncContextManager",
@@ -469,13 +469,23 @@ def __exit__(self, exctype, excinst, exctb):
469469returnFalse
470470471471472+def_lookup_special(obj, name, default):
473+# Follow the standard lookup behaviour for special methods.
474+frominspectimportgetattr_static, _descriptor_get
475+cls=type(obj)
476+try:
477+descr=getattr_static(cls, name)
478+exceptAttributeError:
479+returndefault
480+return_descriptor_get(descr, obj)
481+482+483+_sentinel= ['SENTINEL']
484+485+472486class_BaseExitStack:
473487"""A base class for ExitStack and AsyncExitStack."""
474488475-@staticmethod
476-def_create_exit_wrapper(cm, cm_exit):
477-returnMethodType(cm_exit, cm)
478-479489@staticmethod
480490def_create_cb_wrapper(callback, /, *args, **kwds):
481491def_exit_wrapper(exc_type, exc, tb):
@@ -499,17 +509,8 @@ def push(self, exit):
499509 Also accepts any object with an __exit__ method (registering a call
500510 to the method instead of the object itself).
501511 """
502-# We use an unbound method rather than a bound method to follow
503-# the standard lookup behaviour for special methods.
504-_cb_type=type(exit)
505-506-try:
507-exit_method=_cb_type.__exit__
508-exceptAttributeError:
509-# Not a context manager, so assume it's a callable.
510-self._push_exit_callback(exit)
511-else:
512-self._push_cm_exit(exit, exit_method)
512+exit_method=_lookup_special(exit, '__exit__', exit)
513+self._push_exit_callback(exit_method)
513514returnexit# Allow use as a decorator.
514515515516defenter_context(self, cm):
@@ -518,17 +519,18 @@ def enter_context(self, cm):
518519 If successful, also pushes its __exit__ method as a callback and
519520 returns the result of the __enter__ method.
520521 """
521-# We look up the special methods on the type to match the with
522-# statement.
523-cls=type(cm)
524-try:
525-_enter=cls.__enter__
526-_exit=cls.__exit__
527-exceptAttributeError:
522+_enter=_lookup_special(cm, '__enter__', _sentinel)
523+if_enteris_sentinel:
524+cls=type(cm)
528525raiseTypeError(f"'{cls.__module__}.{cls.__qualname__}' object does "
529-f"not support the context manager protocol") fromNone
530-result=_enter(cm)
531-self._push_cm_exit(cm, _exit)
526+f"not support the context manager protocol")
527+_exit=_lookup_special(cm, '__exit__', _sentinel)
528+if_exitis_sentinel:
529+cls=type(cm)
530+raiseTypeError(f"'{cls.__module__}.{cls.__qualname__}' object does "
531+f"not support the context manager protocol")
532+result=_enter()
533+self._push_exit_callback(_exit)
532534returnresult
533535534536defcallback(self, callback, /, *args, **kwds):
@@ -544,11 +546,6 @@ def callback(self, callback, /, *args, **kwds):
544546self._push_exit_callback(_exit_wrapper)
545547returncallback# Allow use as a decorator
546548547-def_push_cm_exit(self, cm, cm_exit):
548-"""Helper to correctly register callbacks to __exit__ methods."""
549-_exit_wrapper=self._create_exit_wrapper(cm, cm_exit)
550-self._push_exit_callback(_exit_wrapper, True)
551-552549def_push_exit_callback(self, callback, is_sync=True):
553550self._exit_callbacks.append((is_sync, callback))
554551@@ -641,10 +638,6 @@ class AsyncExitStack(_BaseExitStack, AbstractAsyncContextManager):
641638 # connection later in the list raise an exception.
642639 """
643640644-@staticmethod
645-def_create_async_exit_wrapper(cm, cm_exit):
646-returnMethodType(cm_exit, cm)
647-648641@staticmethod
649642def_create_async_cb_wrapper(callback, /, *args, **kwds):
650643asyncdef_exit_wrapper(exc_type, exc, tb):
@@ -657,16 +650,18 @@ async def enter_async_context(self, cm):
657650 If successful, also pushes its __aexit__ method as a callback and
658651 returns the result of the __aenter__ method.
659652 """
660-cls=type(cm)
661-try:
662-_enter=cls.__aenter__
663-_exit=cls.__aexit__
664-exceptAttributeError:
653+_enter=_lookup_special(cm, '__aenter__', _sentinel)
654+if_enteris_sentinel:
655+cls=type(cm)
665656raiseTypeError(f"'{cls.__module__}.{cls.__qualname__}' object does "
666-f"not support the asynchronous context manager protocol"
667- ) fromNone
668-result=await_enter(cm)
669-self._push_async_cm_exit(cm, _exit)
657+f"not support the asynchronous context manager protocol")
658+_exit=_lookup_special(cm, '__aexit__', _sentinel)
659+if_exitis_sentinel:
660+cls=type(cm)
661+raiseTypeError(f"'{cls.__module__}.{cls.__qualname__}' object does "
662+f"not support the asynchronous context manager protocol")
663+result=await_enter()
664+self._push_exit_callback(_exit, False)
670665returnresult
671666672667defpush_async_exit(self, exit):
@@ -677,14 +672,8 @@ def push_async_exit(self, exit):
677672 Also accepts any object with an __aexit__ method (registering a call
678673 to the method instead of the object itself).
679674 """
680-_cb_type=type(exit)
681-try:
682-exit_method=_cb_type.__aexit__
683-exceptAttributeError:
684-# Not an async context manager, so assume it's a coroutine function
685-self._push_exit_callback(exit, False)
686-else:
687-self._push_async_cm_exit(exit, exit_method)
675+exit_method=_lookup_special(exit, '__aexit__', exit)
676+self._push_exit_callback(exit_method, False)
688677returnexit# Allow use as a decorator
689678690679defpush_async_callback(self, callback, /, *args, **kwds):
@@ -704,12 +693,6 @@ async def aclose(self):
704693"""Immediately unwind the context stack."""
705694awaitself.__aexit__(None, None, None)
706695707-def_push_async_cm_exit(self, cm, cm_exit):
708-"""Helper to correctly register coroutine function to __aexit__
709- method."""
710-_exit_wrapper=self._create_async_exit_wrapper(cm, cm_exit)
711-self._push_exit_callback(_exit_wrapper, False)
712-713696asyncdef__aenter__(self):
714697returnself
715698