From aa188fd860a3b3f2169f0cd5dff8a9c5cc40f18c Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Wed, 28 Dec 2016 17:08:21 +0100
Subject: [PATCH] Added basic support for async iter

---
 jinja2/asyncsupport.py |  9 +++++++++
 jinja2/compiler.py     | 16 ++++++++++++++--
 tests/test_async.py    | 11 +++++++++++
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/jinja2/asyncsupport.py b/jinja2/asyncsupport.py
index dc12be9c..0a19b6ae 100644
--- a/jinja2/asyncsupport.py
+++ b/jinja2/asyncsupport.py
@@ -135,3 +135,12 @@ async def auto_await(value):
     if inspect.isawaitable(value):
         return await value
     return value
+
+
+async def auto_iter(iterable):
+    if hasattr(iterable, '__aiter__'):
+        async for item in iterable:
+            yield item
+        return
+    for item in iterable:
+        yield item
diff --git a/jinja2/compiler.py b/jinja2/compiler.py
index 4afda9ad..f2714c67 100644
--- a/jinja2/compiler.py
+++ b/jinja2/compiler.py
@@ -783,7 +783,7 @@ class CodeGenerator(NodeVisitor):
             self.writeline('dummy = lambda *x: None')
 
         if self.environment._async:
-            self.writeline('from jinja2.asyncsupport import auto_await')
+            self.writeline('from jinja2.asyncsupport import auto_await, auto_iter')
 
         # if we want a deferred initialization we cannot move the
         # environment into a local name
@@ -1130,13 +1130,17 @@ class CodeGenerator(NodeVisitor):
                  "loop it's undefined.  Happened in loop on %s" %
                  self.position(node)))
 
-        self.writeline('for ', node)
+        self.writeline(self.environment._async and 'async for ' or 'for ', node)
         self.visit(node.target, loop_frame)
         self.write(extended_loop and ', l_loop in LoopContext(' or ' in ')
 
         # if we have an extened loop and a node test, we filter in the
         # "outer frame".
         if extended_loop and node.test is not None:
+            if self.environment._async:
+                self.fail('loop filters in async mode are currently if the '
+                          'loop uses the special "loop" variable or is '
+                          'recursive.', node.lineno)
             self.write('(')
             self.visit(node.target, loop_frame)
             self.write(' for ')
@@ -1154,7 +1158,11 @@ class CodeGenerator(NodeVisitor):
         elif node.recursive:
             self.write('reciter')
         else:
+            if self.environment._async:
+                self.write('auto_iter(')
             self.visit(node.iter, loop_frame)
+            if self.environment._async:
+                self.write(')')
 
         if node.recursive:
             self.write(', loop_render_func, depth):')
@@ -1194,7 +1202,11 @@ class CodeGenerator(NodeVisitor):
             self.outdent()
             self.start_write(frame, node)
             self.write('loop(')
+            if self.environment._async:
+                self.write('auto_iter(')
             self.visit(node.iter, frame)
+            if self.environment._async:
+                self.write(')')
             self.write(', loop)')
             self.end_write(frame)
 
diff --git a/tests/test_async.py b/tests/test_async.py
index 13af44dd..44463d10 100644
--- a/tests/test_async.py
+++ b/tests/test_async.py
@@ -96,6 +96,17 @@ def test_async_generate():
     assert rv == ['1', '2', '3']
 
 
+@pytest.mark.skipif(not have_async_gen, reason='No async generators')
+def test_async_iteration_in_tmeplates():
+    t = Template('{% for x in rng %}{{ x }}{% endfor %}',
+                 enable_async=True)
+    async def async_iterator():
+        for item in [1, 2, 3]:
+            yield item
+    rv = list(t.generate(rng=async_iterator()))
+    assert rv == ['1', '2', '3']
+
+
 @pytest.fixture
 def test_env_async():
     env = Environment(loader=DictLoader(dict(
-- 
2.47.2