d.py cooldown decorator with async bucket support.

dpy_async_cooldown.py

from typing import Any, Callable, Optional, TypeVar

import discord
from discord.ext import commands


T_contra = TypeVar('T_contra', contravariant=True)


class AsyncCooldownMapping(commands.CooldownMapping[T_contra]):
    async def get_bucket(self, message: T_contra, current: Optional[float] = None) -> Optional[commands.Cooldown]:
        if self._type is commands.BucketType.default:
            return self._cooldown

        self._verify_cache_integrity(current)
        # this change:
        key = await discord.utils.maybe_coroutine(self._bucket_key, message)
        if key not in self._cache:
            bucket = self.create_bucket(message)
            if bucket is not None:
                self._cache[key] = bucket
        else:
            bucket = self._cache[key]

        return bucket

    async def update_rate_limit(
        self, message: T_contra, current: Optional[float] = None, tokens: int = 1
    ) -> Optional[float]:
        bucket = await self.get_bucket(message, current)
        if bucket is None:
            return None
        return bucket.update_rate_limit(current, tokens=tokens)


def async_cooldown(rate: int, per: float, bucket: Callable[[commands.Context], Any]):
    """A cooldown decorator for message commands that allows for coroutines to be
    passed as a cooldown for more complex operations.

    .. note::
        This is registered as a **check**, and does NOT hook into discord.py's native cooldown
        system, so setting :attr:`.commands.Command.cooldown_after_parsing` will NOT affect when
        this cooldown is parsed, but it will be parsed along with the rest of the checks.


    Parameters
    ------------
    rate: :class:`int`
        The number of times a command can be used before triggering a cooldown.
    per: :class:`float`
        The amount of seconds to wait for a cooldown when it's been triggered.
    type: Union[:class:`.commands.BucketType`, Callable[[:class:`.commands.Context`], Any]]
        The type of cooldown to have. If callable, should return a key for the mapping. Can
        be a coroutine.
    """

    mapping = AsyncCooldownMapping(commands.Cooldown(rate, per), bucket)

    async def predicate(ctx):
        cooldown_obj = await mapping.get_bucket(ctx)
        if cooldown_obj:
            retry_after = cooldown_obj.update_rate_limit()
            if retry_after:
                raise commands.CommandOnCooldown(cooldown_obj, retry_after, bucket)  # type: ignore
        return True

    return commands.check(predicate)