Add support for async comprehensions

CHANGES

@@ -11,4 +11,4 @@
 0.4.1   -- Added missing SourceGenerator.visit_arguments()
 0.5     -- Added support for Python 3.5 infix matrix
            multiplication
-0.6     -- Added support for Python 3.6 f-strings.
+0.6     -- Added support for Python 3.6 f-strings and async comprehensions.

astor/code_gen.py

@@ -716,6 +716,8 @@ class SourceGenerator(ExplicitNodeVisitor):
     def visit_comprehension(self, node):
         set_precedence(node, node.iter, *node.ifs)
         set_precedence(Precedence.comprehension_target, node.target)
+        if getattr(node, 'is_async', False):
+            self.write(' async')
         self.write(' for ', node.target, ' in ', node.iter)
         for if_ in node.ifs:
             self.write(' if ', if_)

tests/test_code_gen.py

@@ -271,11 +271,11 @@ class CodegenTestCase(unittest.TestCase):
 
     def test_comprehension(self):
         source = """
-            ((x,y) for x,y in zip(a,b))
+            ((x,y) for x,y in zip(a,b) if a)
         """
         self.assertAstEqual(source)
         source = """
-            fields = [(a, _format(b)) for (a, b) in iter_fields(node)]
+            fields = [(a, _format(b)) for (a, b) in iter_fields(node) if a]
         """
         self.assertAstEqual(source)
         source = """
@@ -284,6 +284,16 @@ class CodegenTestCase(unittest.TestCase):
         """
         self.assertAstEqual(source)
 
+    def test_async_comprehension(self):
+        source = """
+            async def f():
+                [(await x) async for x in y]
+                [(await i) for i in b if await c]
+                (await x async for x in y)
+                {i for i in b async for i in a if await i for b in i}
+        """
+        self.assertAstSourceEqualIfAtLeastVersion(source, (3, 6))
+
     def test_tuple_corner_cases(self):
         source = """
             a = ()