From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Wed, 28 Dec 2016 15:11:09 +0000 (+0100)
Subject: Wrap generate to support async mode
X-Git-Tag: 2.9~66
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=3e6992f271ed6a9855e9a57b54f04d455e5c63b0;p=thirdparty%2Fjinja.git

Wrap generate to support async mode
---

diff --git a/jinja2/asyncsupport.py b/jinja2/asyncsupport.py
index ca897a77..dc12be9c 100644
--- a/jinja2/asyncsupport.py
+++ b/jinja2/asyncsupport.py
@@ -15,6 +15,33 @@ async def concat_async(async_gen):
     return concat(rv)
 
 
+async def generate_async(self, *args, **kwargs):
+    vars = dict(*args, **kwargs)
+    try:
+        async for event in self.root_render_func(self.new_context(vars)):
+            yield event
+    except Exception:
+        exc_info = sys.exc_info()
+    else:
+        return
+    yield self.environment.handle_exception(exc_info, True)
+
+
+def wrap_generate_func(original_generate):
+    def _convert_generator(self, loop, args, kwargs):
+        async_gen = self.generate_async(*args, **kwargs)
+        try:
+            while 1:
+                yield loop.run_until_complete(async_gen.__anext__())
+        except StopAsyncIteration:
+            pass
+    def generate(self, *args, **kwargs):
+        if not self.environment._async:
+            return original_generate(self, *args, **kwargs)
+        return _convert_generator(self, asyncio.get_event_loop(), args, kwargs)
+    return generate
+
+
 async def render_async(self, *args, **kwargs):
     if not self.environment._async:
         raise RuntimeError('The environment was not created with async mode '
@@ -84,6 +111,8 @@ async def make_module_async(self, vars=None, shared=False, locals=None):
 
 def patch_template():
     from jinja2 import Template
+    Template.generate_async = generate_async
+    Template.generate = wrap_generate_func(Template.generate)
     Template.render_async = render_async
     Template.render = wrap_render_func(Template.render)
     Template._get_default_module = wrap_default_module(
diff --git a/tests/test_async.py b/tests/test_async.py
index e21f7413..13af44dd 100644
--- a/tests/test_async.py
+++ b/tests/test_async.py
@@ -88,6 +88,14 @@ def test_async_blocks():
     assert rv == '<Test><Test>'
 
 
+@pytest.mark.skipif(not have_async_gen, reason='No async generators')
+def test_async_generate():
+    t = Template('{% for x in [1, 2, 3] %}{{ x }}{% endfor %}',
+                 enable_async=True)
+    rv = list(t.generate())
+    assert rv == ['1', '2', '3']
+
+
 @pytest.fixture
 def test_env_async():
     env = Environment(loader=DictLoader(dict(