Search code examples
pythondiscord.pypython-decorators

My cooldown decorator isn't functioning properly, how can I fix this?


So I attempted to make my own decorator, very similar to the built-in cooldown decorator provided to us within the discord.py module, except I'm trying to make it work on regular functions (e.g. non "@client.command async def myCommand()" functions) that it doesn't cover.

I have a function that says "hello" back to a user when they say "hello", but I don't want them spamming it over and over and causing the bot to spam it as well. This is what I have currently:

@client.event
async def on_message(message):
    if message.content == "hello":
        try: 
            await sayHello(message)
        except Exception as err:
            print(err)

helloCooldown = commands.CooldownMapping.from_cooldown(1.0, 20.0, BucketType.user)
async def sayHello(message):

    bucket = helloCooldown.get_bucket(message)
    retry_after = bucket.update_rate_limit()
    if retry_after:
        return # exits function if on cooldown / is rate limited

    author = message.author
    channel = message.channel

    await channel.send(f"Hello, {author.name}")

The decorator I created takes the rate, per, and type (similar to the built-in one) thats placed over a non-command function:

def myCooldown(rate, per, type):

    def outer(func):

        def inner(*args):

            cd = commands.CooldownMapping.from_cooldown(rate, per, type)

            bucket = cd.get_bucket(args[0]) # get 'message' argument from the original function

            retry_after = bucket.update_rate_limit()
            if retry_after:
                return # exit out if its on cooldown/ rate limited
            else:
                return func # executes function if not on cooldown

        return inner

    return outer

@myCooldown(1.0, 20.0, BucketType.user)
async def sayHello(message):
    # say hello

The expected behavior is that it stays on cooldown for 20 seconds before saying "hello" again, if it's called. However, I get the error message "Object function cant be used in 'await' expression". How can I fix my decorator to work the way I'm intending it to work?


Solution

  • When you try to await sayHello(message), first sayHello(message) is executed (which is really inner(message)), which returns func.

    Your program tries to await func, which doesn't really make sense, so it throws an error.

    You need to change inner so it returns an await-able object. This means that it can't return None, so you should raise an error instead.

    from discord import DiscordException
    
    class FunctionOnCooldown(DiscordException):
        pass 
    
    def myCooldown(rate, per, type):
        def outer(func):
            cd = commands.CooldownMapping.from_cooldown(rate, per, type)
            def inner(message, *args, **kwargs):
                bucket = cd.get_bucket(message) 
    
                retry_after = bucket.update_rate_limit()
                if retry_after:
                    raise FunctionOnCooldown
                else:
                    return func(*args, **kwargs) # executes function if not on cooldown
    
            return inner
    
        return outer