tuxbot-bot/tuxbot/core/utils/functions/extra.py

51 lines
1.6 KiB
Python
Raw Normal View History

import asyncio
2020-05-24 01:16:08 +02:00
import discord
from discord.ext import commands, flags
2020-05-24 01:16:08 +02:00
class ContextPlus(commands.Context):
async def send(self, content=None, *args, **kwargs):
2020-06-06 18:51:47 +02:00
if (
hasattr(self.command, "deletable") and self.command.deletable
) and kwargs.pop("deletable", True):
message = await super().send(content, *args, **kwargs)
2020-06-06 18:51:47 +02:00
await message.add_reaction("🗑")
def check(reaction: discord.Reaction, user: discord.User):
2020-06-06 18:51:47 +02:00
return (
user == self.author
and str(reaction.emoji) == "🗑"
and reaction.message.id == message.id
)
try:
2020-06-06 18:51:47 +02:00
await self.bot.wait_for("reaction_add", timeout=60.0, check=check)
except asyncio.TimeoutError:
2020-06-06 18:51:47 +02:00
await message.remove_reaction("🗑", self.bot.user)
else:
await message.delete()
return message
else:
return await super().send(content, *args, **kwargs)
class CommandPLus(flags.FlagCommand):
def __init__(self, function, **kwargs):
super().__init__(function, **kwargs)
self.deletable = kwargs.pop("deletable", True)
def command_extra(*args, **kwargs):
return commands.command(*args, **kwargs, cls=CommandPLus)
class GroupPlus(flags.FlagGroup):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.deletable = kwargs.pop("deletable", True)
def group_extra(*args, **kwargs):
return commands.group(*args, **kwargs, cls=GroupPlus)