bpo-44471: Change error type for bad objects in ExitStack.enter_context() (GH-26820)

serhiy-storchaka · web-flow · commit 6cb145d23f5c · 2021-06-29T11:28:15.000+03:00
A TypeError is now raised instead of an AttributeError in
ExitStack.enter_context() and AsyncExitStack.enter_async_context()
for objects which do not support the context manager or
asynchronous context manager protocols correspondingly.
diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst
@@ -515,6 +515,10 @@ Functions and classes provided:
       These context managers may suppress exceptions just as they normally
       would if used directly as part of a :keyword:`with` statement.
 
+      ... versionchanged:: 3.11
+          Raises :exc:`TypeError` instead of :exc:`AttributeError` if *cm*
+          is not a context manager.
+
    .. method:: push(exit)
 
       Adds a context manager's :meth:`__exit__` method to the callback stack.
@@ -585,6 +589,10 @@ Functions and classes provided:
       Similar to :meth:`enter_context` but expects an asynchronous context
       manager.
 
+      ... versionchanged:: 3.11
+          Raises :exc:`TypeError` instead of :exc:`AttributeError` if *cm*
+          is not an asynchronous context manager.
+
    .. method:: push_async_exit(exit)
 
       Similar to :meth:`push` but expects either an asynchronous context manager
diff --git a/Doc/whatsnew/3.11.rst b/Doc/whatsnew/3.11.rst
@@ -75,6 +75,12 @@ New Features
 Other Language Changes
 ======================
 
+A :exc:`TypeError` is now raised instead of an :exc:`AttributeError` in
+:meth:`contextlib.ExitStack.enter_context` and 
+:meth:`contextlib.AsyncExitStack.enter_async_context` for objects which do not
+support the :term:`context manager` or :term:`asynchronous context manager`
+protocols correspondingly.
+(Contributed by Serhiy Storchaka in :issue:`44471`.)
 
 * A :exc:`TypeError` is now raised instead of an :exc:`AttributeError` in 
   :keyword:`with` and :keyword:`async with` statements for objects which do not
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
@@ -473,9 +473,14 @@ def enter_context(self, cm):
         """
         # 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 @@ 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.
         """
-        _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
 
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
@@ -661,6 +661,25 @@ def _exit():
             result.append(2)
         self.assertEqual(result, [1, 2, 3, 4])
 
+    def test_enter_context_errors(self):
+        class LacksEnterAndExit:
+            pass
+        class LacksEnter:
+            def __exit__(self, *exc_info):
+                pass
+        class LacksExit:
+            def __enter__(self):
+                pass
+
+        with self.exit_stack() as stack:
+            with self.assertRaisesRegex(TypeError, 'the context manager'):
+                stack.enter_context(LacksEnterAndExit())
+            with self.assertRaisesRegex(TypeError, 'the context manager'):
+                stack.enter_context(LacksEnter())
+            with self.assertRaisesRegex(TypeError, 'the context manager'):
+                stack.enter_context(LacksExit())
+            self.assertFalse(stack._exit_callbacks)
+
     def test_close(self):
         result = []
         with self.exit_stack() as stack:
@@ -886,9 +905,11 @@ def test_excessive_nesting(self):
     def test_instance_bypass(self):
         class Example(object): pass
         cm = Example()
+        cm.__enter__ = object()
         cm.__exit__ = object()
         stack = self.exit_stack()
-        self.assertRaises(AttributeError, stack.enter_context, cm)
+        with self.assertRaisesRegex(TypeError, 'the context manager'):
+            stack.enter_context(cm)
         stack.push(cm)
         self.assertIs(stack._exit_callbacks[-1][1], cm)
 
diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py
@@ -483,7 +483,7 @@ async def __aexit__(self, *exc_details):
             1/0
 
     @_async_test
-    async def test_async_enter_context(self):
+    async def test_enter_async_context(self):
         class TestCM(object):
             async def __aenter__(self):
                 result.append(1)
@@ -504,6 +504,26 @@ async def _exit():
 
         self.assertEqual(result, [1, 2, 3, 4])
 
+    @_async_test
+    async def test_enter_async_context_errors(self):
+        class LacksEnterAndExit:
+            pass
+        class LacksEnter:
+            async def __aexit__(self, *exc_info):
+                pass
+        class LacksExit:
+            async def __aenter__(self):
+                pass
+
+        async with self.exit_stack() as stack:
+            with self.assertRaisesRegex(TypeError, 'asynchronous context manager'):
+                await stack.enter_async_context(LacksEnterAndExit())
+            with self.assertRaisesRegex(TypeError, 'asynchronous context manager'):
+                await stack.enter_async_context(LacksEnter())
+            with self.assertRaisesRegex(TypeError, 'asynchronous context manager'):
+                await stack.enter_async_context(LacksExit())
+            self.assertFalse(stack._exit_callbacks)
+
     @_async_test
     async def test_async_exit_exception_chaining(self):
         # Ensure exception chaining matches the reference behaviour
@@ -536,6 +556,18 @@ async def suppress_exc(*exc_details):
         self.assertIsInstance(inner_exc, ValueError)
         self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
 
+    @_async_test
+    async def test_instance_bypass_async(self):
+        class Example(object): pass
+        cm = Example()
+        cm.__aenter__ = object()
+        cm.__aexit__ = object()
+        stack = self.exit_stack()
+        with self.assertRaisesRegex(TypeError, 'asynchronous context manager'):
+            await stack.enter_async_context(cm)
+        stack.push_async_exit(cm)
+        self.assertIs(stack._exit_callbacks[-1][1], cm)
+
 
 class TestAsyncNullcontext(unittest.TestCase):
     @_async_test
diff --git a/Misc/NEWS.d/next/Library/2021-06-21-10-46-58.bpo-44471.2QjXv_.rst b/Misc/NEWS.d/next/Library/2021-06-21-10-46-58.bpo-44471.2QjXv_.rst
@@ -0,0 +1,5 @@
+A :exc:`TypeError` is now raised instead of an :exc:`AttributeError` in
+:meth:`contextlib.ExitStack.enter_context` and
+:meth:`contextlib.AsyncExitStack.enter_async_context` for objects which do
+not support the :term:`context manager` or :term:`asynchronous context
+manager` protocols correspondingly.
