Skip to content

Commit

Permalink
avoid an unclosed auto_aiter
Browse files Browse the repository at this point in the history
  • Loading branch information
graingert committed Apr 7, 2024
1 parent 3fd91e4 commit b9fc229
Showing 1 changed file with 20 additions and 5 deletions.
25 changes: 20 additions & 5 deletions src/jinja2/async_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
from .utils import _PassArg
from .utils import pass_eval_context

if t.TYPE_CHECKING:
import typing_extensions as te

V = t.TypeVar("V")


Expand Down Expand Up @@ -67,15 +70,27 @@ async def auto_await(value: t.Union[t.Awaitable["V"], "V"]) -> "V":
return t.cast("V", value)


async def auto_aiter(
class _IteratorToAsyncIterator(t.Generic[V]):
def __init__(self, iterator: "t.Iterator[V]"):
self._iterator = iterator

def __aiter__(self) -> "te.Self":
return self

async def __anext__(self) -> V:
try:
return next(self._iterator)
except StopIteration as e:
raise StopAsyncIteration(e.value) from e


def auto_aiter(
iterable: "t.Union[t.AsyncIterable[V], t.Iterable[V]]",
) -> "t.AsyncIterator[V]":
if hasattr(iterable, "__aiter__"):
async for item in t.cast("t.AsyncIterable[V]", iterable):
yield item
return iterable.__aiter__()
else:
for item in iterable:
yield item
return _IteratorToAsyncIterator(iter(iterable))


async def auto_to_list(
Expand Down

0 comments on commit b9fc229

Please sign in to comment.