aboutsummaryrefslogtreecommitdiffstats
path: root/bot/decorators.py
blob: dfe80e5cda46360824149e09ab16f3706bd22779 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import logging
import random
from asyncio import Lock
from functools import wraps
from weakref import WeakValueDictionary

from discord import Colour, Embed
from discord.ext import commands
from discord.ext.commands import Context

from bot.constants import ERROR_REPLIES

log = logging.getLogger(__name__)


def with_role(*role_ids: int):
    """Check to see whether the invoking user has any of the roles specified in role_ids."""
    async def predicate(ctx: Context):
        if not ctx.guild:  # Return False in a DM
            log.debug(f"{ctx.author} tried to use the '{ctx.command.name}'command from a DM. "
                      "This command is restricted by the with_role decorator. Rejecting request.")
            return False

        for role in ctx.author.roles:
            if role.id in role_ids:
                log.debug(f"{ctx.author} has the '{role.name}' role, and passes the check.")
                return True

        log.debug(f"{ctx.author} does not have the required role to use "
                  f"the '{ctx.command.name}' command, so the request is rejected.")
        return False
    return commands.check(predicate)


def without_role(*role_ids: int):
    """Check whether the invoking user does not have all of the roles specified in role_ids."""
    async def predicate(ctx: Context):
        if not ctx.guild:  # Return False in a DM
            log.debug(f"{ctx.author} tried to use the '{ctx.command.name}' command from a DM. "
                      "This command is restricted by the without_role decorator. Rejecting request.")
            return False

        author_roles = [role.id for role in ctx.author.roles]
        check = all(role not in author_roles for role in role_ids)
        log.debug(f"{ctx.author} tried to call the '{ctx.command.name}' command. "
                  f"The result of the without_role check was {check}.")
        return check
    return commands.check(predicate)


def in_channel(channel_id):
    """Check that the command invocation is in the channel specified by channel_id."""
    async def predicate(ctx: Context):
        check = ctx.channel.id == channel_id
        log.debug(f"{ctx.author} tried to call the '{ctx.command.name}' command. "
                  f"The result of the in_channel check was {check}.")
        return check
    return commands.check(predicate)


def locked():
    """
    Allows the user to only run one instance of the decorated command at a time.

    Subsequent calls to the command from the same author are ignored until the command has completed invocation.

    This decorator has to go before (below) the `command` decorator.
    """
    def wrap(func):
        func.__locks = WeakValueDictionary()

        @wraps(func)
        async def inner(self, ctx, *args, **kwargs):
            lock = func.__locks.setdefault(ctx.author.id, Lock())
            if lock.locked():
                embed = Embed()
                embed.colour = Colour.red()

                log.debug(f"User tried to invoke a locked command.")
                embed.description = (
                    "You're already using this command. Please wait until "
                    "it is done before you use it again."
                )
                embed.title = random.choice(ERROR_REPLIES)
                await ctx.send(embed=embed)
                return

            async with func.__locks.setdefault(ctx.author.id, Lock()):
                return await func(self, ctx, *args, **kwargs)
        return inner
    return wrap