mirror of
https://github.com/Cog-Creators/Red-DiscordBot.git
synced 2025-11-06 11:18:54 -05:00
Make the largest loops lazier and async to avoid blocking (#3767)
* lets reduce config calls here Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * Lets normalize how we name config attributes across the bot. Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * .... Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * Just a tiny PR improving config call in a lot of places (Specially events and Help) Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * stop using `bot.guilds` in `on_command_add` * Just a tiny PR improving config call in a lot of places (Specially events and Help) Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * missed this one Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * nothing to see here Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * lets reduce config calls here Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * Just a tiny PR improving config call in a lot of places (Specially events and Help) Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * stop using `bot.guilds` in `on_command_add` * missed this one Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * welp Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * welp Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * welp Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * jack Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * Update redbot/cogs/mod/kickban.py Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com> * Update redbot/cogs/filter/filter.py Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com> * jack Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * make all large loops async to avoid blocking larger bots Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * ... Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * okay now working AsyncGen Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * may or may not have forgotten black Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * Apply suggestions from code review Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com> * jack's review Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * DOCS Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * DOCS Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * Apply suggestions from code review Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com> * jack Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> * Apply suggestions from code review Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com> * Update redbot/core/utils/__init__.py Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com> * Update redbot/core/utils/__init__.py Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com> * avoid loop if possible and if not only iterate once Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> Co-authored-by: jack1142 <6032823+jack1142@users.noreply.github.com>
This commit is contained in:
Draper
GitHub
parent
465812b673
commit
ad979180e5
@ -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
|
||||
===============
|
||||
|
||||
|
||||
@ -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 = (
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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))
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user