[3.14] GH-135171: Bug fix, test improved

Author: efimov-mikhail

Lib/test/test_asyncgen.py

@@ -1836,11 +1836,18 @@ async def run():
         self.assertEqual(res, [i * 2 for i in range(1, 10)])
 
     def test_async_gen_expression_incorrect(self):
+        async def ag():
+            yield 42
+
+        async def run(arg):
+            (x async for x in arg)
+
         err_msg_async = "'async for' requires an object with " \
             "__aiter__ method, got .*"
 
+        self.loop.run_until_complete(run(ag()))
         with self.assertRaisesRegex(TypeError, err_msg_async):
-            (x async for x in None)
+            self.loop.run_until_complete(run(None))
 
     def test_asyncgen_nonstarted_hooks_are_cancellable(self):
         # See https://bugs.python.org/issue38013

Python/codegen.c

@@ -4758,6 +4758,19 @@ pop_inlined_comprehension_state(compiler *c, location loc,
     return SUCCESS;
 }
 
+static inline int
+codegen_comprehension_iter(compiler *c, comprehension_ty comp)
+{
+    VISIT(c, expr, comp->iter);
+    if (comp->is_async) {
+        ADDOP(c, LOC(comp->iter), GET_AITER);
+    }
+    else {
+        ADDOP(c, LOC(comp->iter), GET_ITER);
+    }
+    return SUCCESS;
+}
+
 static int
 codegen_comprehension(compiler *c, expr_ty e, int type,
                       identifier name, asdl_comprehension_seq *generators, expr_ty elt,
@@ -4777,7 +4790,9 @@ codegen_comprehension(compiler *c, expr_ty e, int type,
 
     outermost = (comprehension_ty) asdl_seq_GET(generators, 0);
     if (is_inlined) {
-        VISIT(c, expr, outermost->iter);
+        if (codegen_comprehension_iter(c, outermost)) {
+            goto error;
+        }
         if (push_inlined_comprehension_state(c, loc, entry, &inline_state)) {
             goto error;
         }
@@ -4851,15 +4866,10 @@ codegen_comprehension(compiler *c, expr_ty e, int type,
     }
     Py_CLEAR(co);
 
-    VISIT(c, expr, outermost->iter);
-    if (type == COMP_GENEXP) {
-        if (outermost->is_async) {
-            ADDOP(c, loc, GET_AITER);
-        }
-        else {
-            ADDOP(c, loc, GET_ITER);
-        }
+    if (codegen_comprehension_iter(c, outermost)) {
+        goto error;
     }
+
     ADDOP_I(c, loc, CALL, 0);
 
     if (is_async_comprehension && type != COMP_GENEXP) {