Cooldowns and max concurrency for application commands (#674)

Co-authored-by: Lala Sabathil <[email protected]>

Author: Dorukyum

discord/commands/commands.py

@@ -26,6 +26,7 @@
 from __future__ import annotations
 
 import asyncio
+import datetime
 import functools
 import inspect
 import re
@@ -111,6 +112,29 @@ class _BaseCommand:
 class ApplicationCommand(_BaseCommand, Generic[CogT, P, T]):
     cog = None
 
+    def __init__(self, func: Callable, **kwargs) -> None:
+        from ..ext.commands.cooldowns import CooldownMapping, BucketType, MaxConcurrency
+
+        try:
+            cooldown = func.__commands_cooldown__
+        except AttributeError:
+            cooldown = kwargs.get('cooldown')
+        
+        if cooldown is None:
+            buckets = CooldownMapping(cooldown, BucketType.default)
+        elif isinstance(cooldown, CooldownMapping):
+            buckets = cooldown
+        else:
+            raise TypeError("Cooldown must be a an instance of CooldownMapping or None.")
+        self._buckets: CooldownMapping = buckets
+
+        try:
+            max_concurrency = func.__commands_max_concurrency__
+        except AttributeError:
+            max_concurrency = kwargs.get('max_concurrency')
+
+        self._max_concurrency: Optional[MaxConcurrency] = max_concurrency
+
     def __repr__(self):
         return f"<discord.commands.{self.__class__.__name__} name={self.name}>"
 
@@ -127,16 +151,48 @@ async def __call__(self, ctx, *args, **kwargs):
         """
         return await self.callback(ctx, *args, **kwargs)
 
+    def _prepare_cooldowns(self, ctx: ApplicationContext):
+        if self._buckets.valid:
+            current = datetime.datetime.now().timestamp()
+            bucket = self._buckets.get_bucket(ctx, current)  # type: ignore (ctx instead of non-existent message)
+
+            if bucket is not None:
+                retry_after = bucket.update_rate_limit(current)
+
+                if retry_after:
+                    from ..ext.commands.errors import CommandOnCooldown
+                    raise CommandOnCooldown(bucket, retry_after, self._buckets.type)  # type: ignore
+
     async def prepare(self, ctx: ApplicationContext) -> None:
         # This should be same across all 3 types
         ctx.command = self
 
         if not await self.can_run(ctx):
             raise CheckFailure(f'The check functions for the command {self.name} failed')
 
-        # TODO: Add cooldown
+        if self._max_concurrency is not None:
+            # For this application, context can be duck-typed as a Message
+            await self._max_concurrency.acquire(ctx)  # type: ignore (ctx instead of non-existent message)
 
-        await self.call_before_hooks(ctx)
+        try:
+            self._prepare_cooldowns(ctx)
+            await self.call_before_hooks(ctx)
+        except:
+            if self._max_concurrency is not None:
+                await self._max_concurrency.release(ctx)  # type: ignore (ctx instead of non-existent message)
+            raise
+
+    def reset_cooldown(self, ctx: ApplicationContext) -> None:
+        """Resets the cooldown on this command.
+
+        Parameters
+        -----------
+        ctx: :class:`.ApplicationContext`
+            The invocation context to reset the cooldown under.
+        """
+        if self._buckets.valid:
+            bucket = self._buckets.get_bucket(ctx)  # type: ignore (ctx instead of non-existent message)
+            bucket.reset()
 
     async def invoke(self, ctx: ApplicationContext) -> None:
         await self.prepare(ctx)
@@ -299,6 +355,10 @@ async def call_after_hooks(self, ctx: ApplicationContext) -> None:
         if hook is not None:
             await hook(ctx)
 
+    @property
+    def cooldown(self):
+        return self._buckets._cooldown
+
     @property
     def full_parent_name(self) -> str:
         """:class:`str`: Retrieves the fully qualified parent command name.
@@ -371,6 +431,11 @@ class SlashCommand(ApplicationCommand):
         :exc:`.CommandError` should be used. Note that if the checks fail then
         :exc:`.CheckFailure` exception is raised to the :func:`.on_application_command_error`
         event.
+    cooldown: Optional[:class:`~discord.ext.commands.Cooldown`]
+        The cooldown applied when the command is invoked. ``None`` if the command
+        doesn't have a cooldown.
+
+        .. versionadded:: 2.0
     """
     type = 1
 
@@ -381,6 +446,7 @@ def __new__(cls, *args, **kwargs) -> SlashCommand:
         return self
 
     def __init__(self, func: Callable, *args, **kwargs) -> None:
+        super().__init__(func, **kwargs)
         if not asyncio.iscoroutinefunction(func):
             raise TypeError("Callback must be a coroutine.")
         self.callback = func
@@ -963,6 +1029,11 @@ class ContextMenuCommand(ApplicationCommand):
         :exc:`.CommandError` should be used. Note that if the checks fail then
         :exc:`.CheckFailure` exception is raised to the :func:`.on_application_command_error`
         event.
+    cooldown: Optional[:class:`~discord.ext.commands.Cooldown`]
+        The cooldown applied when the command is invoked. ``None`` if the command
+        doesn't have a cooldown.
+
+        .. versionadded:: 2.0
     """
     def __new__(cls, *args, **kwargs) -> ContextMenuCommand:
         self = super().__new__(cls)
@@ -971,6 +1042,7 @@ def __new__(cls, *args, **kwargs) -> ContextMenuCommand:
         return self
 
     def __init__(self, func: Callable, *args, **kwargs) -> None:
+        super().__init__(func, **kwargs)
         if not asyncio.iscoroutinefunction(func):
             raise TypeError("Callback must be a coroutine.")
         self.callback = func

discord/ext/commands/core.py

@@ -53,7 +53,13 @@
 from ...errors import *
 from .cooldowns import Cooldown, BucketType, CooldownMapping, MaxConcurrency, DynamicCooldownMapping
 from .converter import run_converters, get_converter, Greedy
-from ...commands import _BaseCommand, slash_command, user_command, message_command
+from ...commands import (
+    ApplicationCommand,
+    _BaseCommand,
+    slash_command,
+    user_command,
+    message_command,
+)
 from .cog import Cog
 from .context import Context
 
@@ -2147,7 +2153,7 @@ def pred(ctx: Context) -> bool:
     return check(pred)
 
 def cooldown(rate: int, per: float, type: Union[BucketType, Callable[[Message], Any]] = BucketType.default) -> Callable[[T], T]:
-    """A decorator that adds a cooldown to a :class:`.Command`
+    """A decorator that adds a cooldown to a command
 
     A cooldown allows a command to only be used a specific amount
     of times in a specific time frame. These cooldowns can be based
@@ -2174,15 +2180,15 @@ def cooldown(rate: int, per: float, type: Union[BucketType, Callable[[Message],
     """
 
     def decorator(func: Union[Command, CoroFunc]) -> Union[Command, CoroFunc]:
-        if isinstance(func, Command):
+        if isinstance(func, (Command, ApplicationCommand)):
             func._buckets = CooldownMapping(Cooldown(rate, per), type)
         else:
             func.__commands_cooldown__ = CooldownMapping(Cooldown(rate, per), type)
         return func
     return decorator  # type: ignore
 
 def dynamic_cooldown(cooldown: Union[BucketType, Callable[[Message], Any]], type: BucketType = BucketType.default) -> Callable[[T], T]:
-    """A decorator that adds a dynamic cooldown to a :class:`.Command`
+    """A decorator that adds a dynamic cooldown to a command
 
     This differs from :func:`.cooldown` in that it takes a function that
     accepts a single parameter of type :class:`.discord.Message` and must
@@ -2222,7 +2228,7 @@ def decorator(func: Union[Command, CoroFunc]) -> Union[Command, CoroFunc]:
     return decorator  # type: ignore
 
 def max_concurrency(number: int, per: BucketType = BucketType.default, *, wait: bool = False) -> Callable[[T], T]:
-    """A decorator that adds a maximum concurrency to a :class:`.Command` or its subclasses.
+    """A decorator that adds a maximum concurrency to a command
 
     This enables you to only allow a certain number of command invocations at the same time,
     for example if a command takes too long or if only one user can use it at a time. This
@@ -2247,7 +2253,7 @@ def max_concurrency(number: int, per: BucketType = BucketType.default, *, wait:
 
     def decorator(func: Union[Command, CoroFunc]) -> Union[Command, CoroFunc]:
         value = MaxConcurrency(number, per=per, wait=wait)
-        if isinstance(func, Command):
+        if isinstance(func, (Command, ApplicationCommand)):
             func._max_concurrency = value
         else:
             func.__commands_max_concurrency__ = value