From f8e428bd5b01fb86032d727a7e8b10e3048b5fcf Mon Sep 17 00:00:00 2001 From: Josh Date: Fri, 4 Dec 2020 09:32:53 +1000 Subject: Add support for chunking AsyncIterator objects --- discord/iterators.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) (limited to 'discord') diff --git a/discord/iterators.py b/discord/iterators.py index 34dcd412..51431651 100644 --- a/discord/iterators.py +++ b/discord/iterators.py @@ -62,6 +62,11 @@ class _AsyncIterator: if ret: return elem + def chunk(self, max_size): + if max_size <= 0: + raise ValueError('async iterator chunk sizes must be greater than 0.') + return _ChunkedAsyncIterator(self, max_size) + def map(self, func): return _MappedAsyncIterator(self, func) @@ -92,6 +97,26 @@ class _AsyncIterator: def _identity(x): return x +class _ChunkedAsyncIterator(_AsyncIterator): + def __init__(self, iterator, max_size): + self.iterator = iterator + self.max_size = max_size + + async def next(self): + ret = [] + n = 0 + while n < self.max_size: + try: + item = await self.iterator.next() + except NoMoreItems: + if ret: + return ret + raise + else: + ret.append(item) + n += 1 + return ret + class _MappedAsyncIterator(_AsyncIterator): def __init__(self, iterator, func): self.iterator = iterator -- cgit v1.2.3