asyncsupport.py 7.7 KB

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