diff options
| -rw-r--r-- | bot/cogs/snekbox.py | 6 | ||||
| -rw-r--r-- | bot/utils/messages.py | 72 | 
2 files changed, 77 insertions, 1 deletions
| diff --git a/bot/cogs/snekbox.py b/bot/cogs/snekbox.py index 2e52b8d1b..184a7545d 100644 --- a/bot/cogs/snekbox.py +++ b/bot/cogs/snekbox.py @@ -12,6 +12,8 @@ from discord.ext.commands import (  from bot.cogs.rmq import RMQ  from bot.constants import Channels, ERROR_REPLIES, NEGATIVE_REPLIES, Roles, URLs +from bot.utils.messages import wait_for_deletion +  log = logging.getLogger(__name__) @@ -181,7 +183,9 @@ class Snekbox:                      else:                          msg = f"{ctx.author.mention} Your eval job has completed.\n\n```py\n{output}\n```" -                    await ctx.send(msg) +                    response = await ctx.send(msg) +                    await wait_for_deletion(response, user_ids=(ctx.author.id,), client=ctx.bot) +                  else:                      await ctx.send(                          f"{ctx.author.mention} Your eval job has completed.\n\n```py\n[No output]\n```" diff --git a/bot/utils/messages.py b/bot/utils/messages.py new file mode 100644 index 000000000..c625beb5c --- /dev/null +++ b/bot/utils/messages.py @@ -0,0 +1,72 @@ +import asyncio +import contextlib +from typing import Sequence + +from discord import Message +from discord.abc import Snowflake + + +async def wait_for_deletion( +    message: Message, +    user_ids: Sequence[Snowflake], +    deletion_emojis: Sequence[str]=("❌",), +    timeout: float=60 * 5, +    attach_emojis=True, +    client=None +): +    """ +    Waits for up to `timeout` seconds for a reaction by +    any of the specified `user_ids` to delete the message. + +    Args: +        message (Message): +            The message that should be monitored for reactions +            and possibly deleted. Must be a message sent on a +            guild since access to the bot instance is required. + +        user_ids (Sequence[Snowflake]): +            A sequence of users that are allowed to delete +            this message. + +    Kwargs: +        deletion_emojis (Sequence[str]): +            A sequence of emojis that are considered deletion +            emojis. + +        timeout (float): +            A positive float denoting the maximum amount of +            time to wait for a deletion reaction. + +        attach_emojis (bool): +            Whether to attach the given `deletion_emojis` +            to the message in the given `context` + +        client (Optional[discord.Client]): +            The client instance handling the original command. +            If not given, will take the client from the guild +            of the message. +    """ + +    if message.guild is None and client is None: +        raise ValueError("Message must be sent on a guild") + +    bot = client or message.guild.me + +    if attach_emojis: +        for emoji in deletion_emojis: +            await message.add_reaction(emoji) + +    def check(reaction, user): +        return ( +            reaction.message.id == message.id and +            reaction.emoji in deletion_emojis and +            user.id in user_ids +        ) + +    with contextlib.suppress(asyncio.TimeoutError): +        await bot.wait_for( +            'reaction_add', +            check=check, +            timeout=timeout +        ) +        await message.delete() | 
