diff --git a/docs/framework_utils.rst b/docs/framework_utils.rst index 055ee72e6..f210bf4d7 100644 --- a/docs/framework_utils.rst +++ b/docs/framework_utils.rst @@ -10,6 +10,17 @@ General Utility .. automodule:: redbot.core.utils :members: deduplicate_iterables, bounded_gather, bounded_gather_iter +.. autoclass:: AsyncIter + :members: + :exclude-members: enumerate, filter + + .. automethod:: enumerate + :async-for: + + .. automethod:: filter + :async-for: + + Chat Formatting =============== diff --git a/redbot/cogs/admin/announcer.py b/redbot/cogs/admin/announcer.py index 3aa420b4b..1485f5109 100644 --- a/redbot/cogs/admin/announcer.py +++ b/redbot/cogs/admin/announcer.py @@ -3,6 +3,7 @@ import asyncio import discord from redbot.core import commands from redbot.core.i18n import Translator +from redbot.core.utils import AsyncIter from redbot.core.utils.chat_formatting import humanize_list, inline _ = Translator("Announcer", __file__) @@ -55,7 +56,7 @@ class Announcer: async def announcer(self): guild_list = self.ctx.bot.guilds failed = [] - for g in guild_list: + async for g in AsyncIter(guild_list, delay=0.5): if not self.active: return @@ -68,7 +69,6 @@ class Announcer: await channel.send(self.message) except discord.Forbidden: failed.append(str(g.id)) - await asyncio.sleep(0.5) if failed: msg = ( diff --git a/redbot/cogs/mod/kickban.py b/redbot/cogs/mod/kickban.py index 0bfe08c9d..697506fdf 100644 --- a/redbot/cogs/mod/kickban.py +++ b/redbot/cogs/mod/kickban.py @@ -7,6 +7,7 @@ from typing import cast, Optional, Union import discord from redbot.core import commands, i18n, checks, modlog +from redbot.core.utils import AsyncIter from redbot.core.utils.chat_formatting import pagify, humanize_number, bold from redbot.core.utils.mod import is_allowed_by_hierarchy, get_audit_reason from .abc import MixinMeta @@ -134,7 +135,7 @@ class KickBanMixin(MixinMeta): async def check_tempban_expirations(self): member = namedtuple("Member", "id guild") while self == self.bot.get_cog("Mod"): - for guild in self.bot.guilds: + async for guild in AsyncIter(self.bot.guilds, steps=100): if not guild.me.guild_permissions.ban_members: continue try: diff --git a/redbot/cogs/permissions/converters.py b/redbot/cogs/permissions/converters.py index cbf1db76d..ab0257eb9 100644 --- a/redbot/cogs/permissions/converters.py +++ b/redbot/cogs/permissions/converters.py @@ -6,6 +6,7 @@ import discord from redbot.core import commands from redbot.core.i18n import Translator +from redbot.core.utils import AsyncIter _ = Translator("PermissionsConverters", __file__) @@ -37,25 +38,25 @@ class GlobalUniqueObjectFinder(commands.Converter): if user is not None: return user - for guild in bot.guilds: + async for guild in AsyncIter(bot.guilds, steps=100): role: discord.Role = guild.get_role(_id) if role is not None: return role - objects = itertools.chain( - bot.get_all_channels(), - bot.users, - bot.guilds, - *(filter(lambda r: not r.is_default(), guild.roles) for guild in bot.guilds), - ) + all_roles = [ + filter(lambda r: not r.is_default(), guild.roles) + async for guild in AsyncIter(bot.guilds, steps=100) + ] + + objects = itertools.chain(bot.get_all_channels(), bot.users, bot.guilds, *all_roles,) maybe_matches = [] - for obj in objects: + async for obj in AsyncIter(objects, steps=100): if obj.name == arg or str(obj) == arg: maybe_matches.append(obj) if ctx.guild is not None: - for member in ctx.guild.members: + async for member in AsyncIter(ctx.guild.members, steps=100): if member.nick == arg and not any(obj.id == member.id for obj in maybe_matches): maybe_matches.append(member) @@ -102,7 +103,7 @@ class GuildUniqueObjectFinder(commands.Converter): ) maybe_matches = [] - for obj in objects: + async for obj in AsyncIter(objects, steps=100): if obj.name == arg or str(obj) == arg: maybe_matches.append(obj) try: diff --git a/redbot/cogs/reports/reports.py b/redbot/cogs/reports/reports.py index fe2971347..ea010b04a 100644 --- a/redbot/cogs/reports/reports.py +++ b/redbot/cogs/reports/reports.py @@ -7,6 +7,7 @@ import contextlib import discord from redbot.core import Config, checks, commands +from redbot.core.utils import AsyncIter from redbot.core.utils.chat_formatting import pagify, box from redbot.core.utils.antispam import AntiSpam from redbot.core.bot import Red @@ -115,7 +116,7 @@ class Reports(commands.Cog): else: perms = discord.Permissions(**permissions) - for guild in self.bot.guilds: + async for guild in AsyncIter(self.bot.guilds, steps=100): x = guild.get_member(author.id) if x is not None: if await self.internal_filter(x, mod, perms): diff --git a/redbot/core/bank.py b/redbot/core/bank.py index 7df78d7ea..baf0a4c8a 100644 --- a/redbot/core/bank.py +++ b/redbot/core/bank.py @@ -12,6 +12,7 @@ from . import Config, errors, commands from .i18n import Translator from .errors import BankPruneError +from .utils import AsyncIter if TYPE_CHECKING: from .bot import Red @@ -405,15 +406,22 @@ async def bank_prune(bot: Red, guild: discord.Guild = None, user_id: int = None) global_bank = await is_global() if global_bank: - _guilds = [g for g in bot.guilds if not g.unavailable and g.large and not g.chunked] - _uguilds = [g for g in bot.guilds if g.unavailable] + _guilds = set() + _uguilds = set() + if user_id is None: + async for g in AsyncIter(bot.guilds, steps=100): + if not g.unavailable and g.large and not g.chunked: + _guilds.add(g) + elif g.unavailable: + _uguilds.add(g) group = _config._get_base_group(_config.USER) else: if guild is None: raise BankPruneError("'guild' can't be None when pruning a local bank") - _guilds = [guild] if not guild.unavailable and guild.large else [] - _uguilds = [guild] if guild.unavailable else [] + if user_id is None: + _guilds = {guild} if not guild.unavailable and guild.large else set() + _uguilds = {guild} if guild.unavailable else set() group = _config._get_base_group(_config.MEMBER, str(guild.id)) if user_id is None: diff --git a/redbot/core/core_commands.py b/redbot/core/core_commands.py index 0e597e8e5..f93ee15e0 100644 --- a/redbot/core/core_commands.py +++ b/redbot/core/core_commands.py @@ -33,6 +33,7 @@ from . import ( i18n, config, ) +from .utils import AsyncIter from .utils.predicates import MessagePredicate from .utils.chat_formatting import ( box, @@ -110,7 +111,7 @@ class CoreLogic: bot._last_exception = exception_log failed_packages.append(name) - for spec, name in cogspecs: + async for spec, name in AsyncIter(cogspecs, steps=10): try: self._cleanup_and_refresh_modules(spec.name) await bot.load_extension(spec) diff --git a/redbot/core/events.py b/redbot/core/events.py index 4622aa149..0fc432a75 100644 --- a/redbot/core/events.py +++ b/redbot/core/events.py @@ -15,6 +15,7 @@ from pkg_resources import DistributionNotFound from redbot.core.commands import RedHelpFormatter, HelpSettings from redbot.core.i18n import Translator +from .utils import AsyncIter from .. import __version__ as red_version, version_info as red_version_info, VersionInfo from . import commands from .config import get_latest_confs @@ -267,7 +268,7 @@ def init_events(bot, cli_flags): if command.qualified_name in disabled_commands: command.enabled = False guild_data = await bot._config.all_guilds() - for guild_id, data in guild_data.items(): + async for guild_id, data in AsyncIter(guild_data.items(), steps=100): disabled_commands = data.get("disabled_commands", []) if command.qualified_name in disabled_commands: command.disable_in(discord.Object(id=guild_id)) diff --git a/redbot/core/utils/__init__.py b/redbot/core/utils/__init__.py index 192320f14..7319626ac 100644 --- a/redbot/core/utils/__init__.py +++ b/redbot/core/utils/__init__.py @@ -1,3 +1,4 @@ +from __future__ import annotations import asyncio import warnings from asyncio import AbstractEventLoop, as_completed, Semaphore @@ -18,9 +19,10 @@ from typing import ( Union, Set, TYPE_CHECKING, + Generator, ) -__all__ = ("bounded_gather", "deduplicate_iterables") +__all__ = ("bounded_gather", "bounded_gather_iter", "deduplicate_iterables", "AsyncIter") _T = TypeVar("_T") @@ -255,3 +257,131 @@ def bounded_gather( tasks = (_sem_wrapper(semaphore, task) for task in coros_or_futures) return asyncio.gather(*tasks, return_exceptions=return_exceptions) + + +class AsyncIter(AsyncIterator[_T], Awaitable[List[_T]]): # pylint: disable=duplicate-bases + """Asynchronous iterator yielding items from ``iterable`` that sleeps for ``delay`` seconds every ``steps`` items. + + Parameters + ---------- + iterable : Iterable + The iterable to make async. + delay: Union[float, int] + The amount of time in seconds to sleep. + steps: int + The number of iterations between sleeps. + """ + + def __init__( + self, iterable: Iterable[_T], delay: Union[float, int] = 0, steps: int = 1 + ) -> None: + self._delay = delay + self._iterator = iter(iterable) + self._i = 0 + self._steps = steps + + def __aiter__(self) -> AsyncIter[_T]: + return self + + async def __anext__(self) -> _T: + try: + item = next(self._iterator) + except StopIteration: + raise StopAsyncIteration + self._i += 1 + if self._i % self._steps == 0: + await asyncio.sleep(self._delay) + return item + + def __await__(self) -> Generator[Any, None, List[_T]]: + return self.flatten().__await__() + + async def flatten(self) -> List[_T]: + """Returns a list of the iterable.""" + return [item async for item in self] + + def filter(self, function: Callable[[_T], Union[bool, Awaitable[bool]]]) -> AsyncFilter[_T]: + """ + Filter the iterable with an (optionally async) predicate. + + Parameters + ---------- + function : Callable[[T], Union[bool, Awaitable[bool]]] + A function or coroutine function which takes one item of ``iterable`` + as an argument, and returns ``True`` or ``False``. + + Returns + ------- + AsyncFilter[T] + An object which can either be awaited to yield a list of the filtered + items, or can also act as an async iterator to yield items one by one. + + Examples + -------- + >>> from redbot.core.utils import AsyncIter + >>> def predicate(value): + ... return value <= 5 + >>> iterator = AsyncIter([1, 10, 5, 100]) + >>> async for i in iterator.filter(predicate): + ... print(i) + 1 + 5 + + >>> from redbot.core.utils import AsyncIter + >>> def predicate(value): + ... return value <= 5 + >>> iterator = AsyncIter([1, 10, 5, 100]) + >>> await iterator.filter(predicate) + [1, 5] + + """ + return async_filter(function, self) + + def enumerate(self, start: int = 0) -> AsyncIterator[Tuple[int, _T]]: + """Async iterable version of `enumerate`. + + Parameters + ---------- + start : int + The index to start from. Defaults to 0. + + Returns + ------- + AsyncIterator[Tuple[int, T]] + An async iterator of tuples in the form of ``(index, item)``. + + Examples + -------- + >>> from redbot.core.utils import AsyncIter + >>> iterator = AsyncIter(['one', 'two', 'three']) + >>> async for i in iterator.enumerate(start=10): + ... print(i) + (10, 'one') + (11, 'two') + (12, 'three') + + """ + return async_enumerate(self, start) + + async def without_duplicates(self) -> AsyncIterator[_T]: + """ + Iterates while omitting duplicated entries. + + Examples + -------- + >>> from redbot.core.utils import AsyncIter + >>> iterator = AsyncIter([1,2,3,3,4,4,5]) + >>> async for i in iterator.without_duplicates(): + ... print(i) + 1 + 2 + 3 + 4 + 5 + """ + _temp = set() + async for item in self: + if item not in _temp: + yield item + _temp.add(item) + del _temp