Add support for chunking AsyncIterator objects

1 files changed, 25 insertions, 0 deletions

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