1# -*- coding: utf-8 -*-
2"""
3    jinja2.asyncsupport
4    ~~~~~~~~~~~~~~~~~~~
5
6    Has all the code for async support which is implemented as a patch
7    for supported Python versions.
8
9    :copyright: (c) 2017 by the Jinja Team.
10    :license: BSD, see LICENSE for more details.
11"""
12import sys
13import asyncio
14import inspect
15from functools import update_wrapper
16
17from jinja2.utils import concat, internalcode, Markup
18from jinja2.environment import TemplateModule
19from jinja2.runtime import LoopContextBase, _last_iteration
20
21
22async def concat_async(async_gen):
23    rv = []
24    async def collect():
25        async for event in async_gen:
26            rv.append(event)
27    await collect()
28    return concat(rv)
29
30
31async def generate_async(self, *args, **kwargs):
32    vars = dict(*args, **kwargs)
33    try:
34        async for event in self.root_render_func(self.new_context(vars)):
35            yield event
36    except Exception:
37        exc_info = sys.exc_info()
38    else:
39        return
40    yield self.environment.handle_exception(exc_info, True)
41
42
43def wrap_generate_func(original_generate):
44    def _convert_generator(self, loop, args, kwargs):
45        async_gen = self.generate_async(*args, **kwargs)
46        try:
47            while 1:
48                yield loop.run_until_complete(async_gen.__anext__())
49        except StopAsyncIteration:
50            pass
51    def generate(self, *args, **kwargs):
52        if not self.environment.is_async:
53            return original_generate(self, *args, **kwargs)
54        return _convert_generator(self, asyncio.get_event_loop(), args, kwargs)
55    return update_wrapper(generate, original_generate)
56
57
58async def render_async(self, *args, **kwargs):
59    if not self.environment.is_async:
60        raise RuntimeError('The environment was not created with async mode '
61                           'enabled.')
62
63    vars = dict(*args, **kwargs)
64    ctx = self.new_context(vars)
65
66    try:
67        return await concat_async(self.root_render_func(ctx))
68    except Exception:
69        exc_info = sys.exc_info()
70    return self.environment.handle_exception(exc_info, True)
71
72
73def wrap_render_func(original_render):
74    def render(self, *args, **kwargs):
75        if not self.environment.is_async:
76            return original_render(self, *args, **kwargs)
77        loop = asyncio.get_event_loop()
78        return loop.run_until_complete(self.render_async(*args, **kwargs))
79    return update_wrapper(render, original_render)
80
81
82def wrap_block_reference_call(original_call):
83    @internalcode
84    async def async_call(self):
85        rv = await concat_async(self._stack[self._depth](self._context))
86        if self._context.eval_ctx.autoescape:
87            rv = Markup(rv)
88        return rv
89
90    @internalcode
91    def __call__(self):
92        if not self._context.environment.is_async:
93            return original_call(self)
94        return async_call(self)
95
96    return update_wrapper(__call__, original_call)
97
98
99def wrap_macro_invoke(original_invoke):
100    @internalcode
101    async def async_invoke(self, arguments, autoescape):
102        rv = await self._func(*arguments)
103        if autoescape:
104            rv = Markup(rv)
105        return rv
106
107    @internalcode
108    def _invoke(self, arguments, autoescape):
109        if not self._environment.is_async:
110            return original_invoke(self, arguments, autoescape)
111        return async_invoke(self, arguments, autoescape)
112    return update_wrapper(_invoke, original_invoke)
113
114
115@internalcode
116async def get_default_module_async(self):
117    if self._module is not None:
118        return self._module
119    self._module = rv = await self.make_module_async()
120    return rv
121
122
123def wrap_default_module(original_default_module):
124    @internalcode
125    def _get_default_module(self):
126        if self.environment.is_async:
127            raise RuntimeError('Template module attribute is unavailable '
128                               'in async mode')
129        return original_default_module(self)
130    return _get_default_module
131
132
133async def make_module_async(self, vars=None, shared=False, locals=None):
134    context = self.new_context(vars, shared, locals)
135    body_stream = []
136    async for item in self.root_render_func(context):
137        body_stream.append(item)
138    return TemplateModule(self, context, body_stream)
139
140
141def patch_template():
142    from jinja2 import Template
143    Template.generate = wrap_generate_func(Template.generate)
144    Template.generate_async = update_wrapper(
145        generate_async, Template.generate_async)
146    Template.render_async = update_wrapper(
147        render_async, Template.render_async)
148    Template.render = wrap_render_func(Template.render)
149    Template._get_default_module = wrap_default_module(
150        Template._get_default_module)
151    Template._get_default_module_async = get_default_module_async
152    Template.make_module_async = update_wrapper(
153        make_module_async, Template.make_module_async)
154
155
156def patch_runtime():
157    from jinja2.runtime import BlockReference, Macro
158    BlockReference.__call__ = wrap_block_reference_call(
159        BlockReference.__call__)
160    Macro._invoke = wrap_macro_invoke(Macro._invoke)
161
162
163def patch_filters():
164    from jinja2.filters import FILTERS
165    from jinja2.asyncfilters import ASYNC_FILTERS
166    FILTERS.update(ASYNC_FILTERS)
167
168
169def patch_all():
170    patch_template()
171    patch_runtime()
172    patch_filters()
173
174
175async def auto_await(value):
176    if inspect.isawaitable(value):
177        return await value
178    return value
179
180
181async def auto_aiter(iterable):
182    if hasattr(iterable, '__aiter__'):
183        async for item in iterable:
184            yield item
185        return
186    for item in iterable:
187        yield item
188
189
190class AsyncLoopContext(LoopContextBase):
191
192    def __init__(self, async_iterator, undefined, after, length, recurse=None,
193                 depth0=0):
194        LoopContextBase.__init__(self, undefined, recurse, depth0)
195        self._async_iterator = async_iterator
196        self._after = after
197        self._length = length
198
199    @property
200    def length(self):
201        if self._length is None:
202            raise TypeError('Loop length for some iterators cannot be '
203                            'lazily calculated in async mode')
204        return self._length
205
206    def __aiter__(self):
207        return AsyncLoopContextIterator(self)
208
209
210class AsyncLoopContextIterator(object):
211    __slots__ = ('context',)
212
213    def __init__(self, context):
214        self.context = context
215
216    def __aiter__(self):
217        return self
218
219    async def __anext__(self):
220        ctx = self.context
221        ctx.index0 += 1
222        if ctx._after is _last_iteration:
223            raise StopAsyncIteration()
224        ctx._before = ctx._current
225        ctx._current = ctx._after
226        try:
227            ctx._after = await ctx._async_iterator.__anext__()
228        except StopAsyncIteration:
229            ctx._after = _last_iteration
230        return ctx._current, ctx
231
232
233async def make_async_loop_context(iterable, undefined, recurse=None, depth0=0):
234    # Length is more complicated and less efficient in async mode.  The
235    # reason for this is that we cannot know if length will be used
236    # upfront but because length is a property we cannot lazily execute it
237    # later.  This means that we need to buffer it up and measure :(
238    #
239    # We however only do this for actual iterators, not for async
240    # iterators as blocking here does not seem like the best idea in the
241    # world.
242    try:
243        length = len(iterable)
244    except (TypeError, AttributeError):
245        if not hasattr(iterable, '__aiter__'):
246            iterable = tuple(iterable)
247            length = len(iterable)
248        else:
249            length = None
250    async_iterator = auto_aiter(iterable)
251    try:
252        after = await async_iterator.__anext__()
253    except StopAsyncIteration:
254        after = _last_iteration
255    return AsyncLoopContext(async_iterator, undefined, after, length, recurse,
256                            depth0)
257