[Utils] Finish and Refactor Predicate Utility (#2169)

* Uses classmethods to create predicates
* Classmethods allow using a combination of different parameters to describe context
* Some predicates assign a captured `result` to the predicate object on success
* Added `ReactionPredicate` equivalent to `MessagePredicate`
* Added `utils.menus.start_adding_reactions`, a non-blocking method for adding reactions asynchronously
* Added documentation
* Uses these new utils throughout the core bot
Happened to also find some bugs in places, and places where we were waiting for events without catching `asyncio.TimeoutError`

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>

redbot/core/commands/context.py

@@ -6,6 +6,7 @@ from discord.ext import commands
 
 from .requires import PermState
 from ..utils.chat_formatting import box
+from ..utils.predicates import MessagePredicate
 from ..utils import common_filters
 
 TICK = "\N{WHITE HEAVY CHECK MARK}"
@@ -141,10 +142,6 @@ class Context(commands.Context):
         messages = tuple(messages)
         ret = []
 
-        more_check = lambda m: (
-            m.author == self.author and m.channel == self.channel and m.content.lower() == "more"
-        )
-
         for idx, page in enumerate(messages, 1):
             if box_lang is None:
                 msg = await self.send(page)
@@ -165,7 +162,11 @@ class Context(commands.Context):
                     "".format(is_are, n_remaining, plural)
                 )
                 try:
-                    resp = await self.bot.wait_for("message", check=more_check, timeout=timeout)
+                    resp = await self.bot.wait_for(
+                        "message",
+                        check=MessagePredicate.lower_equal_to("more", self),
+                        timeout=timeout,
+                    )
                 except asyncio.TimeoutError:
                     await query.delete()
                     break
@@ -175,7 +176,7 @@ class Context(commands.Context):
                     except (discord.HTTPException, AttributeError):
                         # In case the bot can't delete other users' messages,
                         # or is not a bot account
-                        # or chanel is a DM
+                        # or channel is a DM
                         await query.delete()
         return ret
 

redbot/core/core_commands.py

@@ -24,6 +24,7 @@ from redbot.core import __version__
 from redbot.core import checks
 from redbot.core import i18n
 from redbot.core import commands
+from .utils.predicates import MessagePredicate
 from .utils.chat_formatting import pagify, box, inline
 
 if TYPE_CHECKING:
@@ -438,73 +439,63 @@ class Core(commands.Cog, CoreLogic):
     @checks.is_owner()
     async def leave(self, ctx):
         """Leaves server"""
-        author = ctx.author
-        guild = ctx.guild
+        await ctx.send("Are you sure you want me to leave this server? (y/n)")
 
-        await ctx.send("Are you sure you want me to leave this server? Type yes to confirm.")
-
-        def conf_check(m):
-            return m.author == author
-
-        response = await self.bot.wait_for("message", check=conf_check)
-
-        if response.content.lower().strip() == "yes":
-            await ctx.send("Alright. Bye :wave:")
-            log.debug("Leaving '{}'".format(guild.name))
-            await guild.leave()
+        pred = MessagePredicate.yes_or_no(ctx)
+        try:
+            await self.bot.wait_for("message", check=MessagePredicate.yes_or_no(ctx))
+        except asyncio.TimeoutError:
+            await ctx.send("Response timed out.")
+            return
+        else:
+            if pred.result is True:
+                await ctx.send("Alright. Bye :wave:")
+                log.debug("Leaving guild '{}'".format(ctx.guild.name))
+                await ctx.guild.leave()
+            else:
+                await ctx.send("Alright, I'll stay then :)")
 
     @commands.command()
     @checks.is_owner()
     async def servers(self, ctx):
         """Lists and allows to leave servers"""
-        owner = ctx.author
         guilds = sorted(list(self.bot.guilds), key=lambda s: s.name.lower())
         msg = ""
+        responses = []
         for i, server in enumerate(guilds, 1):
             msg += "{}: {}\n".format(i, server.name)
-
-        msg += "\nTo leave a server, just type its number."
+            responses.append(str(i))
 
         for page in pagify(msg, ["\n"]):
             await ctx.send(page)
 
-        def msg_check(m):
-            return m.author == owner
-
-        while msg is not None:
-            try:
-                msg = await self.bot.wait_for("message", check=msg_check, timeout=15)
-            except asyncio.TimeoutError:
-                await ctx.send("I guess not.")
-                break
-            try:
-                msg = int(msg.content) - 1
-                if msg < 0:
-                    break
-                await self.leave_confirmation(guilds[msg], owner, ctx)
-                break
-            except (IndexError, ValueError, AttributeError):
-                pass
-
-    async def leave_confirmation(self, server, owner, ctx):
-        await ctx.send("Are you sure you want me to leave {}? (yes/no)".format(server.name))
-
-        def conf_check(m):
-            return m.author == owner
+        query = await ctx.send("To leave a server, just type its number.")
 
+        pred = MessagePredicate.contained_in(responses, ctx)
         try:
-            msg = await self.bot.wait_for("message", check=conf_check, timeout=15)
-            if msg.content.lower().strip() in ("yes", "y"):
-                if server.owner == ctx.bot.user:
-                    await ctx.send("I cannot leave a guild I am the owner of.")
-                    return
-                await server.leave()
-                if server != ctx.guild:
+            await self.bot.wait_for("message", check=pred, timeout=15)
+        except asyncio.TimeoutError:
+            await query.delete()
+        else:
+            await self.leave_confirmation(guilds[pred.result], ctx)
+
+    async def leave_confirmation(self, guild, ctx):
+        if guild.owner.id == ctx.bot.user.id:
+            await ctx.send("I cannot leave a guild I am the owner of.")
+            return
+
+        await ctx.send("Are you sure you want me to leave {}? (yes/no)".format(guild.name))
+        pred = MessagePredicate.yes_or_no(ctx)
+        try:
+            await self.bot.wait_for("message", check=pred, timeout=15)
+            if pred.result is True:
+                await guild.leave()
+                if guild != ctx.guild:
                     await ctx.send("Done.")
             else:
                 await ctx.send("Alright then.")
         except asyncio.TimeoutError:
-            await ctx.send("I guess not.")
+            await ctx.send("Response timed out.")
 
     @commands.command()
     @checks.is_owner()
@@ -892,10 +883,6 @@ class Core(commands.Cog, CoreLogic):
     @commands.cooldown(1, 60 * 10, commands.BucketType.default)
     async def owner(self, ctx):
         """Sets Red's main owner"""
-
-        def check(m):
-            return m.author == ctx.author and m.channel == ctx.channel
-
         # According to the Python docs this is suitable for cryptographic use
         random = SystemRandom()
         length = random.randint(25, 35)
@@ -919,10 +906,14 @@ class Core(commands.Cog, CoreLogic):
         )
 
         try:
-            message = await ctx.bot.wait_for("message", check=check, timeout=60)
+            message = await ctx.bot.wait_for(
+                "message", check=MessagePredicate.same_context(ctx), timeout=60
+            )
         except asyncio.TimeoutError:
             self.owner.reset_cooldown(ctx)
-            await ctx.send(_("The set owner request has timed out."))
+            await ctx.send(
+                _("The `{prefix}set owner` request has timed out.").format(prefix=ctx.prefix)
+            )
         else:
             if message.content.strip() == token:
                 self.owner.reset_cooldown(ctx)
@@ -1146,18 +1137,20 @@ class Core(commands.Cog, CoreLogic):
             )
             await ctx.send(_("Would you like to receive a copy via DM? (y/n)"))
 
-            def same_author_check(m):
-                return m.author == ctx.author and m.channel == ctx.channel
-
+            pred = MessagePredicate.yes_or_no(ctx)
             try:
-                msg = await ctx.bot.wait_for("message", check=same_author_check, timeout=60)
+                await ctx.bot.wait_for("message", check=pred, timeout=60)
             except asyncio.TimeoutError:
-                await ctx.send(_("Ok then."))
+                await ctx.send(_("Response timed out."))
             else:
-                if msg.content.lower().strip() == "y":
-                    await ctx.author.send(
-                        _("Here's a copy of the backup"), file=discord.File(str(backup_file))
-                    )
+                if pred.result is True:
+                    await ctx.send(_("OK, it's on its way!"))
+                    async with ctx.author.typing():
+                        await ctx.author.send(
+                            _("Here's a copy of the backup"), file=discord.File(str(backup_file))
+                        )
+                else:
+                    await ctx.send(_("OK then."))
         else:
             await ctx.send(_("That directory doesn't seem to exist..."))
 

redbot/core/dev_commands.py

@@ -8,9 +8,11 @@ from contextlib import redirect_stdout
 from copy import copy
 
 import discord
+
 from . import checks, commands
 from .i18n import Translator
 from .utils.chat_formatting import box, pagify
+from .utils.predicates import MessagePredicate
 
 """
 Notice:
@@ -218,12 +220,8 @@ class Dev(commands.Cog):
         self.sessions.add(ctx.channel.id)
         await ctx.send(_("Enter code to execute or evaluate. `exit()` or `quit` to exit."))
 
-        msg_check = lambda m: (
-            m.author == ctx.author and m.channel == ctx.channel and m.content.startswith("`")
-        )
-
         while True:
-            response = await ctx.bot.wait_for("message", check=msg_check)
+            response = await ctx.bot.wait_for("message", check=MessagePredicate.regex(r"^`", ctx))
 
             cleaned = self.cleanup_code(response.content)
 

redbot/core/utils/menus.py

@@ -1,15 +1,14 @@
-"""
-Original source of reaction-based menu idea from
-https://github.com/Lunar-Dust/Dusty-Cogs/blob/master/menu/menu.py
-
-Ported to Red V3 by Palm\_\_ (https://github.com/palmtree5)
-"""
+# Original source of reaction-based menu idea from
+# https://github.com/Lunar-Dust/Dusty-Cogs/blob/master/menu/menu.py
+#
+# Ported to Red V3 by Palm\_\_ (https://github.com/palmtree5)
 import asyncio
 import contextlib
-from typing import Union, Iterable
+from typing import Union, Iterable, Optional
 import discord
 
-from redbot.core import commands
+from .. import commands
+from .predicates import ReactionPredicate
 
 _ReactableEmoji = Union[str, discord.Emoji]
 
@@ -71,18 +70,20 @@ async def menu(
         else:
             message = await ctx.send(current_page)
         # Don't wait for reactions to be added (GH-1797)
-        ctx.bot.loop.create_task(_add_menu_reactions(message, controls.keys()))
+        # noinspection PyAsyncCall
+        start_adding_reactions(message, controls.keys(), ctx.bot.loop)
     else:
         if isinstance(current_page, discord.Embed):
             await message.edit(embed=current_page)
         else:
             await message.edit(content=current_page)
 
-    def react_check(r, u):
-        return u == ctx.author and r.message.id == message.id and str(r.emoji) in controls.keys()
-
     try:
-        react, user = await ctx.bot.wait_for("reaction_add", check=react_check, timeout=timeout)
+        react, user = await ctx.bot.wait_for(
+            "reaction_add",
+            check=ReactionPredicate.with_emojis(tuple(controls.keys()), message, ctx.author),
+            timeout=timeout,
+        )
     except asyncio.TimeoutError:
         try:
             await message.clear_reactions()
@@ -152,12 +153,51 @@ async def close_menu(
     return None
 
 
-async def _add_menu_reactions(message: discord.Message, emojis: Iterable[_ReactableEmoji]):
-    """Add the reactions"""
-    # The task should exit silently if the message is deleted
-    with contextlib.suppress(discord.NotFound):
-        for emoji in emojis:
-            await message.add_reaction(emoji)
+def start_adding_reactions(
+    message: discord.Message,
+    emojis: Iterable[_ReactableEmoji],
+    loop: Optional[asyncio.AbstractEventLoop] = None,
+) -> asyncio.Task:
+    """Start adding reactions to a message.
+
+    This is a non-blocking operation - calling this will schedule the
+    reactions being added, but will the calling code will continue to
+    execute asynchronously. There is no need to await this function.
+
+    This is particularly useful if you wish to start waiting for a
+    reaction whilst the reactions are still being added - in fact,
+    this is exactly what `menu` uses to do that.
+
+    This spawns a `asyncio.Task` object and schedules it on ``loop``.
+    If ``loop`` omitted, the loop will be retreived with
+    `asyncio.get_event_loop`.
+
+    Parameters
+    ----------
+    message: discord.Message
+        The message to add reactions to.
+    emojis : Iterable[Union[str, discord.Emoji]]
+        The emojis to react to the message with.
+    loop : Optional[asyncio.AbstractEventLoop]
+        The event loop.
+
+    Returns
+    -------
+    asyncio.Task
+        The task for the coroutine adding the reactions.
+
+    """
+
+    async def task():
+        # The task should exit silently if the message is deleted
+        with contextlib.suppress(discord.NotFound):
+            for emoji in emojis:
+                await message.add_reaction(emoji)
+
+    if loop is None:
+        loop = asyncio.get_event_loop()
+
+    return loop.create_task(task())
 
 
 DEFAULT_CONTROLS = {"⬅": prev_page, "❌": close_menu, "➡": next_page}