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