async_utils.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. import inspect
  2. import typing as t
  3. from functools import WRAPPER_ASSIGNMENTS
  4. from functools import wraps
  5. from .utils import _PassArg
  6. from .utils import pass_eval_context
  7. if t.TYPE_CHECKING:
  8. import typing_extensions as te
  9. V = t.TypeVar("V")
  10. def async_variant(normal_func): # type: ignore
  11. def decorator(async_func): # type: ignore
  12. pass_arg = _PassArg.from_obj(normal_func)
  13. need_eval_context = pass_arg is None
  14. if pass_arg is _PassArg.environment:
  15. def is_async(args: t.Any) -> bool:
  16. return t.cast(bool, args[0].is_async)
  17. else:
  18. def is_async(args: t.Any) -> bool:
  19. return t.cast(bool, args[0].environment.is_async)
  20. # Take the doc and annotations from the sync function, but the
  21. # name from the async function. Pallets-Sphinx-Themes
  22. # build_function_directive expects __wrapped__ to point to the
  23. # sync function.
  24. async_func_attrs = ("__module__", "__name__", "__qualname__")
  25. normal_func_attrs = tuple(set(WRAPPER_ASSIGNMENTS).difference(async_func_attrs))
  26. @wraps(normal_func, assigned=normal_func_attrs)
  27. @wraps(async_func, assigned=async_func_attrs, updated=())
  28. def wrapper(*args, **kwargs): # type: ignore
  29. b = is_async(args)
  30. if need_eval_context:
  31. args = args[1:]
  32. if b:
  33. return async_func(*args, **kwargs)
  34. return normal_func(*args, **kwargs)
  35. if need_eval_context:
  36. wrapper = pass_eval_context(wrapper)
  37. wrapper.jinja_async_variant = True # type: ignore[attr-defined]
  38. return wrapper
  39. return decorator
  40. _common_primitives = {int, float, bool, str, list, dict, tuple, type(None)}
  41. async def auto_await(value: t.Union[t.Awaitable["V"], "V"]) -> "V":
  42. # Avoid a costly call to isawaitable
  43. if type(value) in _common_primitives:
  44. return t.cast("V", value)
  45. if inspect.isawaitable(value):
  46. return await t.cast("t.Awaitable[V]", value)
  47. return value
  48. class _IteratorToAsyncIterator(t.Generic[V]):
  49. def __init__(self, iterator: "t.Iterator[V]"):
  50. self._iterator = iterator
  51. def __aiter__(self) -> "te.Self":
  52. return self
  53. async def __anext__(self) -> V:
  54. try:
  55. return next(self._iterator)
  56. except StopIteration as e:
  57. raise StopAsyncIteration(e.value) from e
  58. def auto_aiter(
  59. iterable: "t.Union[t.AsyncIterable[V], t.Iterable[V]]",
  60. ) -> "t.AsyncIterator[V]":
  61. if hasattr(iterable, "__aiter__"):
  62. return iterable.__aiter__()
  63. else:
  64. return _IteratorToAsyncIterator(iter(iterable))
  65. async def auto_to_list(
  66. value: "t.Union[t.AsyncIterable[V], t.Iterable[V]]",
  67. ) -> t.List["V"]:
  68. return [x async for x in auto_aiter(value)]