Update test_with.py from 3.14.3 (#7158)

Lib/test/test_coroutines.py

@@ -1250,7 +1250,6 @@ class CoroutineTest(unittest.TestCase):
         with self.assertRaises(ZeroDivisionError):
             run_async(foo())
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: "asynchronous context manager.*__aexit__" does not match "'CM' object does not support the context manager protocol (missed __aexit__ method)"
     def test_with_2(self):
         class CM:
             def __aenter__(self):
@@ -1267,7 +1266,6 @@ class CoroutineTest(unittest.TestCase):
             run_async(foo())
         self.assertIs(body_executed, False)
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: "asynchronous context manager" does not match "'CM' object does not support the context manager protocol (missed __aenter__ method)"
     def test_with_3(self):
         class CM:
             def __aexit__(self):
@@ -1284,7 +1282,6 @@ class CoroutineTest(unittest.TestCase):
             run_async(foo())
         self.assertIs(body_executed, False)
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: "asynchronous context manager" does not match "'CM' object does not support the context manager protocol (missed __aexit__ method)"
     def test_with_4(self):
         class CM:
             pass

Lib/test/test_with.py

@@ -1,9 +1,10 @@
-"""Unit tests for the with statement specified in PEP 343."""
+"""Unit tests for the 'with/async with' statements specified in PEP 343/492."""
 
 
 __author__ = "Mike Bland"
 __email__ = "mbland at acm dot org"
 
+import re
 import sys
 import traceback
 import unittest
@@ -11,6 +12,16 @@ from collections import deque
 from contextlib import _GeneratorContextManager, contextmanager, nullcontext
 
 
+def do_with(obj):
+    with obj:
+        pass
+
+
+async def do_async_with(obj):
+    async with obj:
+        pass
+
+
 class MockContextManager(_GeneratorContextManager):
     def __init__(self, *args):
         super().__init__(*args)
@@ -110,34 +121,77 @@ class FailureTestCase(unittest.TestCase):
             with foo: pass
         self.assertRaises(NameError, fooNotDeclared)
 
-    def testEnterAttributeError1(self):
-        class LacksEnter(object):
-            def __exit__(self, type, value, traceback):
-                pass
+    def testEnterAttributeError(self):
+        class LacksEnter:
+            def __exit__(self, type, value, traceback): ...
 
-        def fooLacksEnter():
-            foo = LacksEnter()
-            with foo: pass
-        self.assertRaisesRegex(TypeError, 'the context manager', fooLacksEnter)
-
-    def testEnterAttributeError2(self):
-        class LacksEnterAndExit(object):
-            pass
-
-        def fooLacksEnterAndExit():
-            foo = LacksEnterAndExit()
-            with foo: pass
-        self.assertRaisesRegex(TypeError, 'the context manager', fooLacksEnterAndExit)
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the context manager protocol "
+            "(missed __enter__ method)"
+        ))):
+            do_with(LacksEnter())
 
     def testExitAttributeError(self):
-        class LacksExit(object):
-            def __enter__(self):
-                pass
+        class LacksExit:
+            def __enter__(self): ...
 
-        def fooLacksExit():
-            foo = LacksExit()
-            with foo: pass
-        self.assertRaisesRegex(TypeError, 'the context manager.*__exit__', fooLacksExit)
+        msg = re.escape((
+            "object does not support the context manager protocol "
+            "(missed __exit__ method)"
+        ))
+        # a missing __exit__ is reported missing before a missing __enter__
+        with self.assertRaisesRegex(TypeError, msg):
+            do_with(object())
+        with self.assertRaisesRegex(TypeError, msg):
+            do_with(LacksExit())
+
+    def testWithForAsyncManager(self):
+        class AsyncManager:
+            async def __aenter__(self): ...
+            async def __aexit__(self, type, value, traceback): ...
+
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the context manager protocol "
+            "(missed __exit__ method) but it supports the asynchronous "
+            "context manager protocol. Did you mean to use 'async with'?"
+        ))):
+            do_with(AsyncManager())
+
+    def testAsyncEnterAttributeError(self):
+        class LacksAsyncEnter:
+            async def __aexit__(self, type, value, traceback): ...
+
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the asynchronous context manager protocol "
+            "(missed __aenter__ method)"
+        ))):
+            do_async_with(LacksAsyncEnter()).send(None)
+
+    def testAsyncExitAttributeError(self):
+        class LacksAsyncExit:
+            async def __aenter__(self): ...
+
+        msg = re.escape((
+            "object does not support the asynchronous context manager protocol "
+            "(missed __aexit__ method)"
+        ))
+        # a missing __aexit__ is reported missing before a missing __aenter__
+        with self.assertRaisesRegex(TypeError, msg):
+            do_async_with(object()).send(None)
+        with self.assertRaisesRegex(TypeError, msg):
+            do_async_with(LacksAsyncExit()).send(None)
+
+    def testAsyncWithForSyncManager(self):
+        class SyncManager:
+            def __enter__(self): ...
+            def __exit__(self, type, value, traceback): ...
+
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the asynchronous context manager protocol "
+            "(missed __aexit__ method) but it supports the context manager "
+            "protocol. Did you mean to use 'with'?"
+        ))):
+            do_async_with(SyncManager()).send(None)
 
     def assertRaisesSyntaxError(self, codestr):
         def shouldRaiseSyntaxError(s):
@@ -190,6 +244,7 @@ class FailureTestCase(unittest.TestCase):
                 pass
         self.assertRaises(RuntimeError, shouldThrow)
 
+
 class ContextmanagerAssertionMixin(object):
 
     def setUp(self):
@@ -624,7 +679,7 @@ class AssignmentTargetTestCase(unittest.TestCase):
         class C: pass
         blah = C()
         with mock_contextmanager_generator() as blah.foo:
-            self.assertEqual(hasattr(blah, "foo"), True)
+            self.assertHasAttr(blah, "foo")
 
     def testMultipleComplexTargets(self):
         class C:
@@ -719,7 +774,7 @@ class NestedWith(unittest.TestCase):
         try:
             with self.Dummy() as a, self.InitRaises():
                 pass
-        except:
+        except RuntimeError:
             pass
         self.assertTrue(a.enter_called)
         self.assertTrue(a.exit_called)
@@ -752,7 +807,7 @@ class NestedWith(unittest.TestCase):
             self.assertEqual(10, b1)
             self.assertEqual(20, b2)
 
-    @unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: 'FrameSummary' object has no attribute 'end_lineno'
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; AttributeError: 'FrameSummary' object has no attribute 'end_lineno'
     def testExceptionLocation(self):
         # The location of an exception raised from
         # __init__, __enter__ or __exit__ of a context

crates/vm/src/frame.rs

@@ -10,7 +10,7 @@ use crate::{
         function::{PyCell, PyCellRef, PyFunction},
         tuple::{PyTuple, PyTupleRef},
     },
-    bytecode::{self, Instruction, LoadAttr, LoadSuperAttr},
+    bytecode::{self, Instruction, LoadAttr, LoadSuperAttr, SpecialMethod},
     convert::{IntoObject, ToPyResult},
     coroutine::Coro,
     exceptions::ExceptionCtor,
@@ -1614,15 +1614,10 @@ impl ExecutingFrame<'_> {
                 // Input: [..., obj]
                 // Output: [..., bound_method]
                 use crate::vm::PyMethod;
-                use bytecode::SpecialMethod;
 
                 let obj = self.pop_value();
-                let method_name = match method.get(arg) {
-                    SpecialMethod::Enter => identifier!(vm, __enter__),
-                    SpecialMethod::Exit => identifier!(vm, __exit__),
-                    SpecialMethod::AEnter => identifier!(vm, __aenter__),
-                    SpecialMethod::AExit => identifier!(vm, __aexit__),
-                };
+                let oparg = method.get(arg);
+                let method_name = get_special_method_name(oparg, vm);
 
                 let bound = match vm.get_special_method(&obj, method_name)? {
                     Some(PyMethod::Function { target, func }) => {
@@ -1633,10 +1628,10 @@ impl ExecutingFrame<'_> {
                     }
                     Some(PyMethod::Attribute(bound)) => bound,
                     None => {
-                        return Err(vm.new_type_error(format!(
-                            "'{}' object does not support the context manager protocol (missed {} method)",
-                            obj.class().name(),
-                            method_name
+                        return Err(vm.new_type_error(get_special_method_error_msg(
+                            oparg,
+                            &obj.class().name(),
+                            special_method_can_suggest(&obj, oparg, vm)?,
                         )));
                     }
                 };
@@ -3451,6 +3446,78 @@ impl fmt::Debug for Frame {
     }
 }
 
+/// _PyEval_SpecialMethodCanSuggest
+fn special_method_can_suggest(
+    obj: &PyObjectRef,
+    oparg: SpecialMethod,
+    vm: &VirtualMachine,
+) -> PyResult<bool> {
+    Ok(match oparg {
+        SpecialMethod::Enter | SpecialMethod::Exit => {
+            vm.get_special_method(obj, get_special_method_name(SpecialMethod::AEnter, vm))?
+                .is_some()
+                && vm
+                    .get_special_method(obj, get_special_method_name(SpecialMethod::AExit, vm))?
+                    .is_some()
+        }
+        SpecialMethod::AEnter | SpecialMethod::AExit => {
+            vm.get_special_method(obj, get_special_method_name(SpecialMethod::Enter, vm))?
+                .is_some()
+                && vm
+                    .get_special_method(obj, get_special_method_name(SpecialMethod::Exit, vm))?
+                    .is_some()
+        }
+    })
+}
+
+fn get_special_method_name(oparg: SpecialMethod, vm: &VirtualMachine) -> &'static PyStrInterned {
+    match oparg {
+        SpecialMethod::Enter => identifier!(vm, __enter__),
+        SpecialMethod::Exit => identifier!(vm, __exit__),
+        SpecialMethod::AEnter => identifier!(vm, __aenter__),
+        SpecialMethod::AExit => identifier!(vm, __aexit__),
+    }
+}
+
+/// _Py_SpecialMethod _Py_SpecialMethods
+fn get_special_method_error_msg(
+    oparg: SpecialMethod,
+    class_name: &str,
+    can_suggest: bool,
+) -> String {
+    if can_suggest {
+        match oparg {
+            SpecialMethod::Enter => format!(
+                "'{class_name}' object does not support the context manager protocol (missed __enter__ method) but it supports the asynchronous context manager protocol. Did you mean to use 'async with'?"
+            ),
+            SpecialMethod::Exit => format!(
+                "'{class_name}' object does not support the context manager protocol (missed __exit__ method) but it supports the asynchronous context manager protocol. Did you mean to use 'async with'?"
+            ),
+            SpecialMethod::AEnter => format!(
+                "'{class_name}' object does not support the asynchronous context manager protocol (missed __aenter__ method) but it supports the context manager protocol. Did you mean to use 'with'?"
+            ),
+            SpecialMethod::AExit => format!(
+                "'{class_name}' object does not support the asynchronous context manager protocol (missed __aexit__ method) but it supports the context manager protocol. Did you mean to use 'with'?"
+            ),
+        }
+    } else {
+        match oparg {
+            SpecialMethod::Enter => format!(
+                "'{class_name}' object does not support the context manager protocol (missed __enter__ method)"
+            ),
+            SpecialMethod::Exit => format!(
+                "'{class_name}' object does not support the context manager protocol (missed __exit__ method)"
+            ),
+            SpecialMethod::AEnter => format!(
+                "'{class_name}' object does not support the asynchronous context manager protocol (missed __aenter__ method)"
+            ),
+            SpecialMethod::AExit => format!(
+                "'{class_name}' object does not support the asynchronous context manager protocol (missed __aexit__ method)"
+            ),
+        }
+    }
+}
+
 fn is_module_initializing(module: &PyObject, vm: &VirtualMachine) -> bool {
     let Ok(spec) = module.get_attr(&vm.ctx.new_str("__spec__"), vm) else {
         return false;