fix breakage from async generator

asottile · Jul 15, 2021 · d064832 · d064832
1 parent 4eba9f3
commit d064832
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/pyupgrade/_plugins/generator_expressions_pep289.py b/pyupgrade/_plugins/generator_expressions_pep289.py
@@ -51,7 +51,11 @@ def visit_Call(
             isinstance(node.func, ast.Name) and
             node.func.id in ALLOWED_FUNCS and
             node.args and
-            isinstance(node.args[0], ast.ListComp)
+            isinstance(node.args[0], ast.ListComp) and
+            not any(
+                generator.is_async
+                for generator in node.args[0].generators
+            )
     ):
         if len(node.args) == 1 and not node.keywords:
             yield ast_to_offset(node.args[0]), _delete_list_comp_brackets

diff --git a/tests/features/generator_expressions_pep289_test.py b/tests/features/generator_expressions_pep289_test.py
@@ -20,6 +20,14 @@
             '"".join([[i for _ in range(2)] for i in range(3)])\n',
             id='string join (left alone for perf reasons)',
         ),
+        pytest.param(
+            'async def foo():\n'
+            '    for i in range(3):\n'
+            '        yield i\n'
+            'async def bar():\n'
+            '    sum([i async for i in foo()])\n',
+            id='Contains async',
+        ),
     ),
 )
 def test_fix_generator_expressions_noop(s):