From a68f2f0578bbf812fa2264d0e0bb388340d6e230 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Tue, 3 Apr 2018 01:41:38 +0300
Subject: [PATCH] bpo-29922: Improve error messages in 'async with' (GH-6352)

when __aenter__() or __aexit__() return non-awaitable object.
---
 Lib/test/test_coroutines.py                   | 13 +++++++---
 .../2018-04-03-00-30-25.bpo-29922.CdLuMl.rst  |  2 ++
 Python/ceval.c                                | 25 +++++++++++++++++++
 3 files changed, 36 insertions(+), 4 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2018-04-03-00-30-25.bpo-29922.CdLuMl.rst

diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
index fe26199f95a..9ad7ff9564b 100644
--- a/Lib/test/test_coroutines.py
+++ b/Lib/test/test_coroutines.py
@@ -1255,7 +1255,9 @@ class CoroutineTest(unittest.TestCase):
                 pass
 
         with self.assertRaisesRegex(
-            TypeError, "object int can't be used in 'await' expression"):
+                TypeError,
+                "'async with' received an object from __aenter__ "
+                "that does not implement __await__: int"):
             # it's important that __aexit__ wasn't called
             run_async(foo())
 
@@ -1275,7 +1277,9 @@ class CoroutineTest(unittest.TestCase):
             run_async(foo())
         except TypeError as exc:
             self.assertRegex(
-                exc.args[0], "object int can't be used in 'await' expression")
+                exc.args[0],
+                "'async with' received an object from __aexit__ "
+                "that does not implement __await__: int")
             self.assertTrue(exc.__context__ is not None)
             self.assertTrue(isinstance(exc.__context__, ZeroDivisionError))
         else:
@@ -1299,8 +1303,9 @@ class CoroutineTest(unittest.TestCase):
 
 
         with self.assertRaisesRegex(
-            TypeError, "object int can't be used in 'await' expression"):
-
+                TypeError,
+                "'async with' received an object from __aexit__ "
+                "that does not implement __await__: int"):
             run_async(foo())
 
         self.assertEqual(CNT, 1)
diff --git a/Misc/NEWS.d/next/Core and Builtins/2018-04-03-00-30-25.bpo-29922.CdLuMl.rst b/Misc/NEWS.d/next/Core and Builtins/2018-04-03-00-30-25.bpo-29922.CdLuMl.rst
new file mode 100644
index 00000000000..d8c144e59d6
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2018-04-03-00-30-25.bpo-29922.CdLuMl.rst	
@@ -0,0 +1,2 @@
+Improved error messages in 'async with' when ``__aenter__()`` or
+``__aexit__()`` return non-awaitable object.
diff --git a/Python/ceval.c b/Python/ceval.c
index d18a284e9f9..da83e41f5dd 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -69,6 +69,7 @@ static PyObject * unicode_concatenate(PyObject *, PyObject *,
 static PyObject * special_lookup(PyObject *, _Py_Identifier *);
 static int check_args_iterable(PyObject *func, PyObject *vararg);
 static void format_kwargs_mapping_error(PyObject *func, PyObject *kwargs);
+static void format_awaitable_error(PyTypeObject *, int);
 
 #define NAME_ERROR_MSG \
     "name '%.200s' is not defined"
@@ -1736,6 +1737,11 @@ main_loop:
             PyObject *iterable = TOP();
             PyObject *iter = _PyCoro_GetAwaitableIter(iterable);
 
+            if (iter == NULL) {
+                format_awaitable_error(Py_TYPE(iterable),
+                                       _Py_OPCODE(next_instr[-2]));
+            }
+
             Py_DECREF(iterable);
 
             if (iter != NULL && PyCoro_CheckExact(iter)) {
@@ -4985,6 +4991,25 @@ format_exc_unbound(PyCodeObject *co, int oparg)
     }
 }
 
+static void
+format_awaitable_error(PyTypeObject *type, int prevopcode)
+{
+    if (type->tp_as_async == NULL || type->tp_as_async->am_await == NULL) {
+        if (prevopcode == BEFORE_ASYNC_WITH) {
+            PyErr_Format(PyExc_TypeError,
+                         "'async with' received an object from __aenter__ "
+                         "that does not implement __await__: %.100s",
+                         type->tp_name);
+        }
+        else if (prevopcode == WITH_CLEANUP_START) {
+            PyErr_Format(PyExc_TypeError,
+                         "'async with' received an object from __aexit__ "
+                         "that does not implement __await__: %.100s",
+                         type->tp_name);
+        }
+    }
+}
+
 static PyObject *
 unicode_concatenate(PyObject *v, PyObject *w,
                     PyFrameObject *f, const _Py_CODEUNIT *next_instr)