aioitertools icon indicating copy to clipboard operation
aioitertools copied to clipboard

Peeking into an iterator

Open wagnerflo opened this issue 5 years ago • 0 comments

Hypothetical code, but not very far from my use case:

from chardet.universaldetector import UniversalDetector

with peekerator(bytestream) as (peek, bytestream):
    detector = UniversalDetector()

    async for chunk in peek:
        detector.feed(chunk)
        if detector.done:
            break

    encoding = detector.close().get('encoding')
    async for line in iterdecode(bytestream, encoding):
        print(line)

With bytestream being something like starlette.requests.Request.stream(). And here's my implementation of peekerator, which I think would be a nice addition to aioitertools.

from collections import deque
from contextlib import contextmanager
from aioitertools.builtins import iter

@contextmanager
def peekerator(iterator):
    it = iter(iterator)
    cache = deque()

    async def peek():
        async for el in it:
            cache.append(el)
            yield el

    async def real():
        while True:
            if cache:
                yield cache.popleft()
            else:
                try:
                    yield await it.__anext__()
                except StopAsyncIteration:
                    return

    yield peek(),real()

The idea is based on more_itertools.peekable but with an peeking iterator instead of method and without the added overhead of implementing prepend and indexing.

If there's interest in adding this, I'd package it up with tests and documentation as a pull request. The name is, of course, up for discussion.

wagnerflo avatar Dec 22 '20 18:12 wagnerflo