From 3c2674725a861a8db62d41d76941df68b95db9a0 Mon Sep 17 00:00:00 2001
From: Josh <josh.ja.butt@gmail.com>
Date: Mon, 26 Apr 2021 13:36:03 +1000
Subject: [PATCH] Add as_chunks helper function
---
discord/utils.py | 70 ++++++++++++++++++++++++++++++++++++++++++++++++
docs/api.rst | 2 ++
2 files changed, 72 insertions(+)
diff --git a/discord/utils.py b/discord/utils.py
index 49d3f4413..88948da57 100644
--- a/discord/utils.py
+++ b/discord/utils.py
@@ -28,6 +28,7 @@ import asyncio
import collections.abc
from typing import (
Any,
+ AsyncIterator,
Callable,
Dict,
Generic,
@@ -67,6 +68,7 @@ __all__ = (
'remove_markdown',
'escape_markdown',
'escape_mentions',
+ 'as_chunks',
)
DISCORD_EPOCH = 1420070400000
@@ -103,6 +105,7 @@ else:
T = TypeVar('T')
T_co = TypeVar('T_co', covariant=True)
+_Iter = Union[Iterator[T], AsyncIterator[T]]
CSP = TypeVar('CSP', bound='CachedSlotProperty')
@@ -723,3 +726,70 @@ def escape_mentions(text: str) -> str:
The text with the mentions removed.
"""
return re.sub(r'@(everyone|here|[!&]?[0-9]{17,20})', '@\u200b\\1', text)
+
+
+def _chunk(iterator: Iterator[T], max_size: int) -> Iterator[List[T]]:
+ ret = []
+ n = 0
+ for item in iterator:
+ ret.append(item)
+ n += 1
+ if n == max_size:
+ yield ret
+ ret = []
+ n = 0
+ if ret:
+ yield ret
+
+async def _achunk(iterator: AsyncIterator[T], max_size: int) -> AsyncIterator[List[T]]:
+ ret = []
+ n = 0
+ async for item in iterator:
+ ret.append(item)
+ n += 1
+ if n == max_size:
+ yield ret
+ ret = []
+ n = 0
+ if ret:
+ yield ret
+
+
+@overload
+def as_chunks(iterator: Iterator[T], max_size: int) -> Iterator[List[T]]:
+ ...
+
+
+@overload
+def as_chunks(iterator: AsyncIterator[T], max_size: int) -> AsyncIterator[List[T]]:
+ ...
+
+
+def as_chunks(iterator: _Iter[T], max_size: int) -> _Iter[List[T]]:
+ """A helper function that collects an iterator into chunks of a given size.
+
+ .. versionadded:: 2.0
+
+ Parameters
+ ----------
+ iterator: Union[:class:`collections.abc.Iterator`, :class:`collections.abc.AsyncIterator`]
+ The iterator to chunk, can be sync or async.
+ max_size: :class:`int`
+ The maximum chunk size.
+
+
+ .. warning::
+
+ The last chunk collected may not be as large as ``max_size``.
+
+ Returns
+ --------
+ Union[:class:`Iterator`, :class:`AsyncIterator`]
+ A new iterator which yields chunks of a given size.
+ """
+ if max_size <= 0:
+ raise ValueError('Chunk sizes must be greater than 0.')
+
+ if isinstance(iterator, AsyncIterator):
+ return _achunk(iterator, max_size)
+ return _chunk(iterator, max_size)
diff --git a/docs/api.rst b/docs/api.rst
index a1ceb5b33..c2f5950a2 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -925,6 +925,8 @@ Utility Functions
.. autofunction:: discord.utils.utcnow
+.. autofunction:: discord.utils.as_chunks
+
.. _discord-api-enums:
Enumerations