[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>
This commit is contained in:
Toby Harradine
2018-10-06 08:07:09 +10:00
committed by GitHub
parent 5d44bfabed
commit dea9dde637
15 changed files with 1229 additions and 320 deletions

View File

@@ -13,8 +13,16 @@ import time
import redbot.core
from redbot.core import Config, commands, checks, bank
from redbot.core.data_manager import cog_data_path
from redbot.core.utils.menus import menu, DEFAULT_CONTROLS, prev_page, next_page, close_menu
from redbot.core.i18n import Translator, cog_i18n
from redbot.core.utils.menus import (
menu,
DEFAULT_CONTROLS,
prev_page,
next_page,
close_menu,
start_adding_reactions,
)
from redbot.core.utils.predicates import MessagePredicate, ReactionPredicate
from urllib.parse import urlparse
from .manager import shutdown_lavalink_server
@@ -225,22 +233,17 @@ class Audio(commands.Cog):
async def dj(self, ctx):
"""Toggle DJ mode (users need a role to use audio commands)."""
dj_role_id = await self.config.guild(ctx.guild).dj_role()
if dj_role_id is None:
if dj_role_id is None and ctx.guild.get_role(dj_role_id):
await self._embed_msg(
ctx, "Please set a role to use with DJ mode. Enter the role name now."
ctx, "Please set a role to use with DJ mode. Enter the role name or ID now."
)
def check(m):
return m.author == ctx.author
try:
dj_role = await ctx.bot.wait_for("message", timeout=15.0, check=check)
dj_role_obj = discord.utils.get(ctx.guild.roles, name=dj_role.content)
if dj_role_obj is None:
return await self._embed_msg(ctx, "No role with that name.")
await ctx.invoke(self.role, dj_role_obj)
pred = MessagePredicate.valid_role(ctx)
await ctx.bot.wait_for("message", timeout=15.0, check=pred)
await ctx.invoke(self.role, pred.result)
except asyncio.TimeoutError:
return await self._embed_msg(ctx, "No role entered, try again later.")
return await self._embed_msg(ctx, "Response timed out, try again later.")
dj_enabled = await self.config.guild(ctx.guild).dj_enabled()
await self.config.guild(ctx.guild).dj_enabled.set(not dj_enabled)
@@ -710,20 +713,21 @@ class Audio(commands.Cog):
return
if player.current:
for i in range(4):
await message.add_reaction(expected[i])
def check(r, u):
return (
r.message.id == message.id
and u == ctx.message.author
and any(e in str(r.emoji) for e in expected)
)
task = start_adding_reactions(message, expected[:4], ctx.bot.loop)
else:
task = None
try:
(r, u) = await self.bot.wait_for("reaction_add", check=check, timeout=10.0)
(r, u) = await self.bot.wait_for(
"reaction_add",
check=ReactionPredicate.with_emojis(expected, message, ctx.author),
timeout=10.0,
)
except asyncio.TimeoutError:
return await self._clear_react(message)
else:
if task is not None:
task.cancel()
reacts = {v: k for k, v in emoji.items()}
react = reacts[r.emoji]
if react == "prev":
@@ -1125,11 +1129,12 @@ class Audio(commands.Cog):
if not playlist_name:
await self._embed_msg(ctx, "Please enter a name for this playlist.")
def check(m):
return m.author == ctx.author and not m.content.startswith(ctx.prefix)
try:
playlist_name_msg = await ctx.bot.wait_for("message", timeout=15.0, check=check)
playlist_name_msg = await ctx.bot.wait_for(
"message",
timeout=15.0,
check=MessagePredicate.regex(fr"^(?!{ctx.prefix})", ctx),
)
playlist_name = playlist_name_msg.content.split(" ")[0].strip('"')
if len(playlist_name) > 20:
return await self._embed_msg(ctx, "Try the command again with a shorter name.")
@@ -1238,11 +1243,10 @@ class Audio(commands.Cog):
ctx, "Please upload the playlist file. Any other message will cancel this operation."
)
def check(m):
return m.author == ctx.author
try:
file_message = await ctx.bot.wait_for("message", timeout=30.0, check=check)
file_message = await ctx.bot.wait_for(
"message", timeout=30.0, check=MessagePredicate.same_context(ctx)
)
except asyncio.TimeoutError:
return await self._embed_msg(ctx, "No file detected, try again later.")
try:

View File

@@ -9,6 +9,7 @@ from redbot.core.bot import Red
from redbot.core.i18n import Translator, cog_i18n
from redbot.core.utils.mod import slow_deletion, mass_purge
from redbot.cogs.mod.log import log
from redbot.core.utils.predicates import MessagePredicate
_ = Translator("Cleanup", __file__)
@@ -31,13 +32,10 @@ class Cleanup(commands.Cog):
Tries its best to cleanup after itself if the response is positive.
"""
def author_check(message):
return message.author == ctx.author
prompt = await ctx.send(
_("Are you sure you want to delete {} messages? (y/n)").format(number)
)
response = await ctx.bot.wait_for("message", check=author_check)
response = await ctx.bot.wait_for("message", check=MessagePredicate.same_context(ctx))
if response.content.lower().startswith("y"):
await prompt.delete()

View File

@@ -11,6 +11,7 @@ import discord
from redbot.core import Config, checks, commands
from redbot.core.utils.chat_formatting import box, pagify
from redbot.core.i18n import Translator, cog_i18n
from redbot.core.utils.predicates import MessagePredicate
_ = Translator("CustomCommands", __file__)
@@ -58,14 +59,11 @@ class CommandObj:
).format("customcommand", "customcommand", "exit()")
await ctx.send(intro)
def check(m):
return m.channel == ctx.channel and m.author == ctx.message.author
responses = []
args = None
while True:
await ctx.send(_("Add a random response:"))
msg = await self.bot.wait_for("message", check=check)
msg = await self.bot.wait_for("message", check=MessagePredicate.same_context(ctx))
if msg.content.lower() == "exit()":
break
@@ -130,18 +128,27 @@ class CommandObj:
author = ctx.message.author
def check(m):
return m.channel == ctx.channel and m.author == ctx.author
if ask_for and not response:
await ctx.send(_("Do you want to create a 'randomized' cc? {}").format("y/n"))
await ctx.send(_("Do you want to create a 'randomized' custom command? (y/n)"))
msg = await self.bot.wait_for("message", check=check)
if msg.content.lower() == "y":
pred = MessagePredicate.yes_or_no(ctx)
try:
await self.bot.wait_for("message", check=pred, timeout=30)
except TimeoutError:
await ctx.send(_("Response timed out, please try again later."))
return
if pred.result is True:
response = await self.get_responses(ctx=ctx)
else:
await ctx.send(_("What response do you want?"))
response = (await self.bot.wait_for("message", check=check)).content
try:
resp = await self.bot.wait_for(
"message", check=MessagePredicate.same_context(ctx), timeout=180
)
except TimeoutError:
await ctx.send(_("Response timed out, please try again later."))
return
response = resp.content
if response:
# test to raise

View File

@@ -6,6 +6,7 @@ from redbot.core.bot import Red
from redbot.core.i18n import Translator, cog_i18n
from redbot.cogs.dataconverter.core_specs import SpecResolver
from redbot.core.utils.chat_formatting import box
from redbot.core.utils.predicates import MessagePredicate
_ = Translator("DataConverter", __file__)
@@ -48,11 +49,10 @@ class DataConverter(commands.Cog):
menu_message = await ctx.send(box(menu))
def pred(m):
return m.channel == ctx.channel and m.author == ctx.author
try:
message = await self.bot.wait_for("message", check=pred, timeout=60)
message = await self.bot.wait_for(
"message", check=MessagePredicate.same_context(ctx), timeout=60
)
except asyncio.TimeoutError:
return await ctx.send(_("Try this again when you are more ready"))
else:

View File

@@ -1,7 +1,7 @@
import asyncio
import discord
from redbot.core import commands
from redbot.core.utils.predicates import MessagePredicate
__all__ = ["do_install_agreement"]
@@ -21,13 +21,12 @@ async def do_install_agreement(ctx: commands.Context):
if downloader is None or downloader.already_agreed:
return True
def does_agree(msg: discord.Message):
return ctx.author == msg.author and ctx.channel == msg.channel and msg.content == "I agree"
await ctx.send(REPO_INSTALL_MSG)
try:
await ctx.bot.wait_for("message", check=does_agree, timeout=30)
await ctx.bot.wait_for(
"message", check=MessagePredicate.lower_equal_to("i agree", ctx), timeout=30
)
except asyncio.TimeoutError:
await ctx.send("Your response has timed out, please try again.")
return False

View File

@@ -11,6 +11,8 @@ from redbot.core import checks, commands, config
from redbot.core.bot import Red
from redbot.core.i18n import Translator, cog_i18n
from redbot.core.utils.chat_formatting import box
from redbot.core.utils.menus import start_adding_reactions
from redbot.core.utils.predicates import ReactionPredicate, MessagePredicate
from .converters import CogOrCommand, RuleType, ClearableRuleType
@@ -20,9 +22,6 @@ COG = "COG"
COMMAND = "COMMAND"
GLOBAL = 0
# noinspection PyDictDuplicateKeys
REACTS = {"\N{WHITE HEAVY CHECK MARK}": True, "\N{NEGATIVE SQUARED CROSS MARK}": False}
Y_OR_N = {"y": True, "yes": True, "n": False, "no": False}
# The strings in the schema are constants and should get extracted, but not translated until
# runtime.
translate = _
@@ -566,35 +565,29 @@ class Permissions(commands.Cog):
"""Ask "Are you sure?" and get the response as a bool."""
if ctx.guild is None or ctx.guild.me.permissions_in(ctx.channel).add_reactions:
msg = await ctx.send(_("Are you sure?"))
for emoji in REACTS.keys():
await msg.add_reaction(emoji)
# noinspection PyAsyncCall
task = start_adding_reactions(msg, ReactionPredicate.YES_OR_NO_EMOJIS, ctx.bot.loop)
pred = ReactionPredicate.yes_or_no(msg, ctx.author)
try:
reaction, user = await ctx.bot.wait_for(
"reaction_add",
check=lambda r, u: (
r.message.id == msg.id and u == ctx.author and r.emoji in REACTS
),
timeout=30,
)
await ctx.bot.wait_for("reaction_add", check=pred, timeout=30)
except asyncio.TimeoutError:
agreed = False
await ctx.send(_("Response timed out."))
return False
else:
agreed = REACTS.get(reaction.emoji)
task.cancel()
agreed = pred.result
finally:
await msg.delete()
else:
await ctx.send(_("Are you sure? (y/n)"))
pred = MessagePredicate.yes_or_no(ctx)
try:
message = await ctx.bot.wait_for(
"message",
check=lambda m: m.author == ctx.author
and m.channel == ctx.channel
and m.content in Y_OR_N,
timeout=30,
)
await ctx.bot.wait_for("message", check=pred, timeout=30)
except asyncio.TimeoutError:
agreed = False
await ctx.send(_("Response timed out."))
return False
else:
agreed = Y_OR_N.get(message.content.lower())
agreed = pred.result
if agreed is False:
await ctx.send(_("Action cancelled."))

View File

@@ -11,6 +11,7 @@ from redbot.core.utils.chat_formatting import pagify, box
from redbot.core.utils.antispam import AntiSpam
from redbot.core.bot import Red
from redbot.core.i18n import Translator, cog_i18n
from redbot.core.utils.predicates import MessagePredicate
from redbot.core.utils.tunnel import Tunnel
@@ -136,13 +137,14 @@ class Reports(commands.Cog):
output += "\n{}".format(prompt)
for page in pagify(output, delims=["\n"]):
dm = await author.send(box(page))
def pred(m):
return m.author == author and m.channel == dm.channel
await author.send(box(page))
try:
message = await self.bot.wait_for("message", check=pred, timeout=45)
message = await self.bot.wait_for(
"message",
check=MessagePredicate.same_context(channel=author.dm_channel, user=author),
timeout=45,
)
except asyncio.TimeoutError:
await author.send(_("You took too long to select. Try again later."))
return None
@@ -247,7 +249,7 @@ class Reports(commands.Cog):
val = await self.send_report(_m, guild)
else:
try:
dm = await author.send(
await author.send(
_(
"Please respond to this message with your Report."
"\nYour report should be a single message"
@@ -256,11 +258,12 @@ class Reports(commands.Cog):
except discord.Forbidden:
return await ctx.send(_("This requires DMs enabled."))
def pred(m):
return m.author == author and m.channel == dm.channel
try:
message = await self.bot.wait_for("message", check=pred, timeout=180)
message = await self.bot.wait_for(
"message",
check=MessagePredicate.same_context(ctx, channel=author.dm_channel),
timeout=180,
)
except asyncio.TimeoutError:
return await author.send(_("You took too long. Try again later."))
else:

View File

@@ -5,6 +5,7 @@ import discord
from redbot.core import Config, checks, commands
from redbot.core.i18n import Translator
from redbot.core.utils.predicates import MessagePredicate
_ = Translator("Warnings", __file__)
@@ -95,11 +96,10 @@ async def get_command_for_exceeded_points(ctx: commands.Context):
await ctx.send(_("You may enter your response now."))
def same_author_check(m):
return m.author == ctx.author
try:
msg = await ctx.bot.wait_for("message", check=same_author_check, timeout=30)
msg = await ctx.bot.wait_for(
"message", check=MessagePredicate.same_context(ctx), timeout=30
)
except asyncio.TimeoutError:
return None
else:
@@ -140,11 +140,10 @@ async def get_command_for_dropping_points(ctx: commands.Context):
await ctx.send(_("You may enter your response now."))
def same_author_check(m):
return m.author == ctx.author
try:
msg = await ctx.bot.wait_for("message", check=same_author_check, timeout=30)
msg = await ctx.bot.wait_for(
"message", check=MessagePredicate.same_context(ctx), timeout=30
)
except asyncio.TimeoutError:
return None
else:

View File

@@ -15,6 +15,7 @@ from redbot.core.i18n import Translator, cog_i18n
from redbot.core.utils.mod import is_admin_or_superior
from redbot.core.utils.chat_formatting import warning, pagify
from redbot.core.utils.menus import menu, DEFAULT_CONTROLS
from redbot.core.utils.predicates import MessagePredicate
_ = Translator("Warnings", __file__)
@@ -363,12 +364,11 @@ class Warnings(commands.Cog):
"""Handles getting description and points for custom reasons"""
to_add = {"points": 0, "description": ""}
def same_author_check(m):
return m.author == ctx.author
await ctx.send(_("How many points should be given for this reason?"))
try:
msg = await ctx.bot.wait_for("message", check=same_author_check, timeout=30)
msg = await ctx.bot.wait_for(
"message", check=MessagePredicate.same_context(ctx), timeout=30
)
except asyncio.TimeoutError:
await ctx.send(_("Ok then."))
return
@@ -385,7 +385,9 @@ class Warnings(commands.Cog):
await ctx.send(_("Enter a description for this reason."))
try:
msg = await ctx.bot.wait_for("message", check=same_author_check, timeout=30)
msg = await ctx.bot.wait_for(
"message", check=MessagePredicate.same_context(ctx), timeout=30
)
except asyncio.TimeoutError:
await ctx.send(_("Ok then."))
return