Match async get iter code generation

DinoV · DinoV · commit 00f9478edbbf · 2021-01-28T11:14:38.000-08:00
Summary: 3.7 no longer emits this random LOAD_CONST/YIELD_FROM pair - this is beacuse Python 3.7 drops support for asynchronous __aiter__: https://bugs.python.org/issue31709 Test Plan: ./python -m test.test_compiler
diff --git a/compiler/pycodegen.py b/compiler/pycodegen.py
@@ -555,6 +555,10 @@ def visitFor(self, node):
             self.visit(node.orelse)
         self.nextBlock(after)
 
+    def emitAsyncIterYieldFrom(self):
+        self.emit('LOAD_CONST', None)
+        self.emit('YIELD_FROM')
+
     def visitAsyncFor(self, node):
         try_ = self.newBlock('async_for_try')
         except_ = self.newBlock('except')
@@ -570,8 +574,7 @@ def visitAsyncFor(self, node):
 
         self.visit(node.iter)
         self.emit('GET_AITER')
-        self.emit('LOAD_CONST', None)
-        self.emit('YIELD_FROM')
+        self.emitAsyncIterYieldFrom()
 
         self.nextBlock(try_)
 
@@ -801,8 +804,7 @@ class Holder: pass
         self.visit(node.generators[0].iter)
         if node.generators[0].is_async:
             self.emit('GET_AITER')
-            self.emit('LOAD_CONST', None)
-            self.emit('YIELD_FROM')
+            self.emitAsyncIterYieldFrom()
         else:
             self.emit('GET_ITER')
         self.emit('CALL_FUNCTION', 1)
@@ -844,8 +846,7 @@ def compile_async_comprehension(self, generators, gen_index, elt, val, type):
         else:
             self.visit(gen.iter)
             self.emit('GET_AITER')
-            self.emit('LOAD_CONST', None)
-            self.emit('YIELD_FROM')
+            self.emitAsyncIterYieldFrom()
 
         self.nextBlock(try_)
         self.emit('SETUP_EXCEPT', except_)
@@ -2081,6 +2082,15 @@ def visitConstant(self, node: ast.Constant):
         self.update_lineno(node)
         self.emit('LOAD_CONST', node.value)
 
+    def emitAsyncIterYieldFrom(self):
+        pass
+
+    def visitAsyncWith(self, node):
+        if not self.scope.coroutine:
+            raise SyntaxError("'async with' outside async function", self.syntax_error_position(node))
+
+        return super().visitAsyncWith(node)
+
 
 def get_default_generator():
     if sys.version_info >= (3, 7):
diff --git a/test_compiler/test_errors.py b/test_compiler/test_errors.py
@@ -1,6 +1,7 @@
 import ast
 import dis
 import inspect
+import sys
 import unittest
 from unittest import TestCase
 from .common import CompilerTest
@@ -51,6 +52,21 @@ def test_return_outside_func_class(self):
                 class C:
                     return""")
 
+    @unittest.skipIf(sys.version_info < (3, 7), "parse fails on 3.6")
+    def test_async_with_outside_async_func(self):
+        with self.assertRaisesRegex(SyntaxError, "'async with' outside async function"):
+            self.compile("""
+                def f():
+                    async with x:
+                        pass""")
+
+    @unittest.skipIf(sys.version_info < (3, 7), "parse fails on 3.6")
+    def test_async_with_outside_func(self):
+        with self.assertRaisesRegex(SyntaxError, "'async with' outside async function"):
+            self.compile("""
+                async with x:
+                    pass""")
+
     def test_return_with_value_async_gen(self):
         with self.assertRaisesRegex(
             SyntaxError, "'return' with value in async generator"
diff --git a/test_compiler/test_py37.py b/test_compiler/test_py37.py
@@ -1,6 +1,7 @@
 import ast
 import dis
 from .common import CompilerTest
+from compiler.pyassem import PyFlowGraph
 from compiler.pycodegen import CodeGenerator, Python37CodeGenerator
 from compiler.optimizer import AstOptimizer
 from compiler.unparse import to_expr
@@ -60,6 +61,28 @@ def test_future_annotations_flag(self):
             CO_NOFREE |  CO_FUTURE_ANNOTATIONS,
         )
 
+    def test_async_aiter(self):
+        # Make sure GET_AITER isn't followed by LOAD_CONST None as Python 3.7 doesn't support async __aiter__
+        for generator in (Python37CodeGenerator, CodeGenerator):
+            outer_graph = self.to_graph("""
+                async def f():
+                    async for x in y:
+                        pass
+            """, generator)
+            for outer_instr in self.graph_to_instrs(outer_graph):
+                if outer_instr.opname == "LOAD_CONST" and isinstance(outer_instr.oparg, CodeGenerator):
+                    saw_aiter = False
+                    for instr in self.graph_to_instrs(outer_instr.oparg.graph):
+                        if saw_aiter:
+                            self.assertEqual(instr.opname == "LOAD_CONST", generator is CodeGenerator)
+                            if generator is CodeGenerator:
+                                self.assertIsNone(instr.oparg)
+                            break
+
+                        if instr.opname == "GET_AITER":
+                            saw_aiter = True
+                    break
+
     def test_future_annotations(self):
         annotations = ["42"]
         for annotation in annotations:
