improve error messages when incorrectly using context managers

python · picnixz · Apr 19, 2025 · Apr 7, 2025 · Apr 13, 2025 · Apr 13, 2025
commit aff8184ad8c56a3edb5da4f8a1ed9ad41353a5ce
diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h
@@ -279,6 +279,7 @@ PyAPI_DATA(const conversion_func) _PyEval_ConversionFuncs[];
 typedef struct _special_method {
     PyObject *name;
     const char *error;
+    const char *error_suggestion;  // improved optional suggestion
 } _Py_SpecialMethod;
 
 PyAPI_DATA(const _Py_SpecialMethod) _Py_SpecialMethods[];
@@ -309,6 +310,8 @@ PyAPI_FUNC(PyObject *) _PyEval_LoadName(PyThreadState *tstate, _PyInterpreterFra
 PyAPI_FUNC(int)
 _Py_Check_ArgsIterable(PyThreadState *tstate, PyObject *func, PyObject *args);
 
+PyAPI_FUNC(int) _PyEval_SpecialMethodCanSuggest(PyObject *self, int oparg);
+
 /* Bits that can be set in PyThreadState.eval_breaker */
 #define _PY_GIL_DROP_REQUEST_BIT (1U << 0)
 #define _PY_SIGNALS_PENDING_BIT (1U << 1)

diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
@@ -5,6 +5,7 @@
 import math
 import opcode
 import os
+import re
 import unittest
 import sys
 import ast
@@ -24,6 +25,35 @@
 from test.support.bytecode_helper import instructions_with_positions
 from test.support.os_helper import FakePath
 
+
+class DummyEnter:
+    def __enter__(self, *args, **kwargs):
+        pass
+
+
+class DummyExit:
+    def __exit__(self, *args, **kwargs):
+        pass
+
+
+class SyncDummy(DummyEnter, DummyExit):
+    pass
+
+
+class AsyncDummyEnter:
+    async def __aenter__(self, *args, **kwargs):
+        pass
+
+
+class AsyncDummyExit:
+    async def __aexit__(self, *args, **kwargs):
+        pass
+
+
+class AsyncDummy(AsyncDummyEnter, AsyncDummyExit):
+    pass
+
+
 class TestSpecifics(unittest.TestCase):
 
     def compile_single(self, source):
@@ -1636,6 +1666,69 @@ async def name_4():
                     pass
             [[]]
 
+    def test_invalid_with_usages(self):
+        def f(obj):
+            with obj:
+                pass
+
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the context manager protocol "
+            "(missed __exit__ method)"
+        ))):
+            f(DummyEnter())
+
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the context manager protocol "
+            "(missed __enter__ method)"
+        ))):
+            f(DummyExit())
+
+        # a missing __exit__ is reported missing before a missing __enter__
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the context manager protocol "
+            "(missed __exit__ method)"
+        ))):
+            f(object())
+
+        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'?"
+        ))):
            f(AsyncDummy())
+
+    def test_invalid_async_with_usages(self):
+        async def f(obj):
+            async with obj:
+                pass
+
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the asynchronous context manager protocol "
+            "(missed __aexit__ method)"
+        ))):
+            f(AsyncDummyEnter()).send(None)
+
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the asynchronous context manager protocol "
+            "(missed __aenter__ method)"
+        ))):
+            f(AsyncDummyExit()).send(None)
+
+        # a missing __aexit__ is reported missing before a missing __aenter__
+        with self.assertRaisesRegex(TypeError, re.escape((
+            "object does not support the asynchronous context manager protocol "
+            "(missed __aexit__ method)"
+        ))):
+            f(object()).send(None)
+
+        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'?"
+        ))):
+            f(SyncDummy()).send(None)
+
+
 class TestBooleanExpression(unittest.TestCase):
     class Value:
         def __init__(self):

diff --git a/Lib/unittest/async_case.py b/Lib/unittest/async_case.py
@@ -75,9 +75,17 @@ async def enterAsyncContext(self, cm):
             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
+            msg = (f"'{cls.__module__}.{cls.__qualname__}' object does "
+                   "not support the asynchronous context manager protocol")
+            try:
+                cls.__enter__
+                cls.__exit__
+            except AttributeError:
+                pass
+            else:
+                msg += (" but it supports the context manager protocol. "
+                        "Did you mean to use enterContext()?")
+            raise TypeError(msg) from None
         result = await enter(cm)
         self.addAsyncCleanup(exit, cm, None, None, None)
         return result

diff --git a/Lib/unittest/case.py b/Lib/unittest/case.py
@@ -111,8 +111,17 @@ def _enter_context(cm, addcleanup):
         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
+        msg = (f"'{cls.__module__}.{cls.__qualname__}' object does "
+               "not support the context manager protocol")
+        try:
+            cls.__aenter__
+            cls.__aexit__
+        except AttributeError:
+            pass
+        else:
+            msg += (" but it supports the asynchronous context manager "
+                    "protocol. Did you mean to use enterAsyncContext()?")
+        raise TypeError(msg) from None
     result = enter(cm)
     addcleanup(exit, cm, None, None, None)
     return result

diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2025-04-07-13-46-57.gh-issue-128398.gJ2zIF.rst b/Misc/NEWS.d/next/Core_and_Builtins/2025-04-07-13-46-57.gh-issue-128398.gJ2zIF.rst
@@ -0,0 +1,4 @@
+Improve error message when an object supporting the synchronous (resp.
+asynchronous) context manager protocol is entered using :keyword:`async
+with` (resp. :keyword:`with`) block instead of :keyword:`with` (resp.
+:keyword:`async with`). Patch by Bénédikt Tran.
@@ -3373,9 +3373,12 @@ dummy_func(
             PyObject *attr_o = _PyObject_LookupSpecialMethod(owner_o, name, &self_or_null_o);
             if (attr_o == NULL) {
                 if (!_PyErr_Occurred(tstate)) {
-                    _PyErr_Format(tstate, PyExc_TypeError,
-                                  _Py_SpecialMethods[oparg].error,
-                                  Py_TYPE(owner_o)->tp_name);
+                    const char *errfmt = _PyEval_SpecialMethodCanSuggest(owner_o, oparg)
+                        ? _Py_SpecialMethods[oparg].error_suggestion
+                        : _Py_SpecialMethods[oparg].error;
+                    assert(!_PyErr_Occurred(tstate));
+                    assert(errfmt != NULL);
+                    _PyErr_Format(tstate, PyExc_TypeError, errfmt, owner_o);
                 }
                 ERROR_IF(true, error);
             }

@@ -545,23 +545,51 @@ const conversion_func _PyEval_ConversionFuncs[4] = {
 const _Py_SpecialMethod _Py_SpecialMethods[] = {
     [SPECIAL___ENTER__] = {
         .name = &_Py_ID(__enter__),
-        .error = "'%.200s' object does not support the "
-                 "context manager protocol (missed __enter__ method)",
+        .error = (
+            "'%T' object does not support the context manager protocol "
+            "(missed __enter__ method)"
+        ),
+        .error_suggestion = (
+            "'%T' 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'?"
+        )
     },
     [SPECIAL___EXIT__] = {
         .name = &_Py_ID(__exit__),
-        .error = "'%.200s' object does not support the "
-                 "context manager protocol (missed __exit__ method)",
+        .error = (
+            "'%T' object does not support the context manager protocol "
+            "(missed __exit__ method)"
+        ),
+        .error_suggestion = (
+            "'%T' 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'?"
+        )
     },
     [SPECIAL___AENTER__] = {
         .name = &_Py_ID(__aenter__),
-        .error = "'%.200s' object does not support the asynchronous "
-                 "context manager protocol (missed __aenter__ method)",
+        .error = (
+            "'%T' object does not support the asynchronous "
+            "context manager protocol (missed __aenter__ method)"
+        ),
+        .error_suggestion = (
+            "'%T' 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'?"
+        )
     },
     [SPECIAL___AEXIT__] = {
         .name = &_Py_ID(__aexit__),
-        .error = "'%.200s' object does not support the asynchronous "
-                 "context manager protocol (missed __aexit__ method)",
+        .error = (
+            "'%T' object does not support the asynchronous "
+            "context manager protocol (missed __aexit__ method)"
+        ),
+        .error_suggestion = (
+            "'%T' 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'?"
+        )
     }
 };
 
@@ -3363,3 +3391,33 @@ _PyEval_LoadName(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObject *na
     }
     return value;
 }
+
+/* Check if a 'cls' provides the given special method. */
+static inline int
+type_has_special_method(PyTypeObject *cls, PyObject *name)
+{
+    // _PyType_Lookup() does not set an exception and returns a borrowed ref
+    assert(!PyErr_Occurred());
+    PyObject *r = _PyType_Lookup(cls, name);
+    return r != NULL && Py_TYPE(r)->tp_descr_get != NULL;
+}
+
+int
+_PyEval_SpecialMethodCanSuggest(PyObject *self, int oparg)
+{
+    PyTypeObject *type = Py_TYPE(self);
+    switch (oparg) {
+        case SPECIAL___ENTER__:
+        case SPEC
B51C
IAL___EXIT__: {
+            return type_has_special_method(type, &_Py_ID(__aenter__))
+                   && type_has_special_method(type, &_Py_ID(__aexit__));
+        }
+        case SPECIAL___AENTER__:
+        case SPECIAL___AEXIT__: {
+            return type_has_special_method(type, &_Py_ID(__enter__))
+                   && type_has_special_method(type, &_Py_ID(__exit__));
+        }
+        default:
+            Py_FatalError("unsupported special method");
+    }
+}
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h