fix(typing): support (Interaction) -> Any as type in @cooldown

disnake/ext/commands/base_core.py

@@ -8,15 +8,7 @@
 import inspect
 from abc import ABC
 from collections.abc import Callable
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    TypeAlias,
-    TypedDict,
-    TypeVar,
-    cast,
-    overload,
-)
+from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict, TypeVar, overload
 
 from disnake import utils
 from disnake.app_commands import ApplicationCommand
@@ -55,7 +47,7 @@ class _AppCommandArgs(TypedDict, total=False):
         guild_only: bool
         extras: dict[str, Any] | None
         checks: list[AppCheck]
-        cooldown: CooldownMapping | None
+        cooldown: CooldownMapping[ApplicationCommandInteraction] | None
         max_concurrency: MaxConcurrency | None
 
 
@@ -197,7 +189,7 @@ def __init__(
         else:
             msg = "Cooldown must be a an instance of CooldownMapping or None."
             raise TypeError(msg)
-        self._buckets: CooldownMapping = buckets
+        self._buckets: CooldownMapping[ApplicationCommandInteraction] = buckets
 
         try:
             max_concurrency = func.__commands_max_concurrency__
@@ -222,7 +214,8 @@ def _ensure_assignment_on_copy(self, other: AppCommandT) -> AppCommandT:
             other._buckets = self._buckets.copy()
         if self._max_concurrency != other._max_concurrency:
             # _max_concurrency won't be None at this point
-            other._max_concurrency = cast("MaxConcurrency", self._max_concurrency).copy()
+            assert self._max_concurrency is not None
+            other._max_concurrency = self._max_concurrency.copy()
 
         if (
             # see https://github.com/DisnakeDev/disnake/pull/678#discussion_r938113624:
@@ -387,11 +380,11 @@ def _prepare_cooldowns(self, inter: ApplicationCommandInteraction) -> None:
         if self._buckets.valid:
             dt = inter.created_at
             current = dt.replace(tzinfo=datetime.timezone.utc).timestamp()
-            bucket = self._buckets.get_bucket(inter, current)  # pyright: ignore[reportArgumentType]
+            bucket = self._buckets.get_bucket(inter, current)
             if bucket is not None:  # pyright: ignore[reportUnnecessaryComparison]
                 retry_after = bucket.update_rate_limit(current)
                 if retry_after:
-                    raise CommandOnCooldown(bucket, retry_after, self._buckets.type)  # pyright: ignore[reportArgumentType]
+                    raise CommandOnCooldown(bucket, retry_after, self._buckets.type)
 
     async def prepare(self, inter: ApplicationCommandInteraction) -> None:
         inter.application_command = self
@@ -401,14 +394,14 @@ async def prepare(self, inter: ApplicationCommandInteraction) -> None:
             raise CheckFailure(msg)
 
         if self._max_concurrency is not None:
-            await self._max_concurrency.acquire(inter)  # pyright: ignore[reportArgumentType]
+            await self._max_concurrency.acquire(inter)
 
         try:
             self._prepare_cooldowns(inter)
             await self.call_before_hooks(inter)
         except Exception:
             if self._max_concurrency is not None:
-                await self._max_concurrency.release(inter)  # pyright: ignore[reportArgumentType]
+                await self._max_concurrency.release(inter)
             raise
 
     def is_on_cooldown(self, inter: ApplicationCommandInteraction) -> bool:
@@ -427,7 +420,7 @@ def is_on_cooldown(self, inter: ApplicationCommandInteraction) -> bool:
         if not self._buckets.valid:
             return False
 
-        bucket = self._buckets.get_bucket(inter)  # pyright: ignore[reportArgumentType]
+        bucket = self._buckets.get_bucket(inter)
         dt = inter.created_at
         current = dt.replace(tzinfo=datetime.timezone.utc).timestamp()
         return bucket.get_tokens(current) == 0
@@ -441,7 +434,7 @@ def reset_cooldown(self, inter: ApplicationCommandInteraction) -> None:
             The interaction with this application command
         """
         if self._buckets.valid:
-            bucket = self._buckets.get_bucket(inter)  # pyright: ignore[reportArgumentType]
+            bucket = self._buckets.get_bucket(inter)
             bucket.reset()
 
     def get_cooldown_retry_after(self, inter: ApplicationCommandInteraction) -> float:
@@ -459,7 +452,7 @@ def get_cooldown_retry_after(self, inter: ApplicationCommandInteraction) -> floa
             If this is ``0.0`` then the command isn't on cooldown.
         """
         if self._buckets.valid:
-            bucket = self._buckets.get_bucket(inter)  # pyright: ignore[reportArgumentType]
+            bucket = self._buckets.get_bucket(inter)
             dt = inter.created_at
             current = dt.replace(tzinfo=datetime.timezone.utc).timestamp()
             return bucket.get_retry_after(current)
@@ -483,7 +476,7 @@ async def invoke(self, inter: ApplicationCommandInteraction, *args: Any, **kwarg
             raise CommandInvokeError(exc) from exc
         finally:
             if self._max_concurrency is not None:
-                await self._max_concurrency.release(inter)  # pyright: ignore[reportArgumentType]
+                await self._max_concurrency.release(inter)
 
             await self.call_after_hooks(inter)
 

disnake/ext/commands/cooldowns.py

@@ -6,8 +6,11 @@
 import time
 from collections import deque
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Generic, Protocol
 
+from typing_extensions import TypeVar
+
+import disnake
 from disnake.enums import Enum
 from disnake.member import Member
 
@@ -16,8 +19,6 @@
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    from ...message import Message
-
 __all__ = (
     "BucketType",
     "Cooldown",
@@ -27,6 +28,19 @@
 )
 
 
+# Context, Interaction, or Message
+class ContextType(Protocol):
+    @property
+    def author(self) -> disnake.abc.User: ...
+    @property
+    def guild(self) -> disnake.Guild | None: ...
+    @property
+    def channel(self) -> disnake.abc.MessageableChannel: ...
+
+
+ContextTypeT = TypeVar("ContextTypeT", bound=ContextType, infer_variance=True)
+
+
 class BucketType(Enum):
     """Specifies a type of bucket for, e.g. a cooldown."""
 
@@ -48,7 +62,7 @@ class BucketType(Enum):
     .. versionadded:: 1.3
     """
 
-    def get_key(self, msg: Message) -> Any:
+    def get_key(self, msg: ContextType) -> Any:
         if self is BucketType.user:
             return msg.author.id
         elif self is BucketType.guild:
@@ -69,7 +83,7 @@ def get_key(self, msg: Message) -> Any:
             ).id
         return None
 
-    def __call__(self, msg: Message) -> Any:
+    def __call__(self, msg: ContextType) -> Any:
         return self.get_key(msg)
 
 
@@ -86,7 +100,7 @@ class Cooldown:
 
     __slots__ = ("rate", "per", "_window", "_tokens", "_last")
 
-    def __init__(self, rate: float, per: float) -> None:
+    def __init__(self, rate: int, per: float) -> None:
         self.rate: int = int(rate)
         self.per: float = float(per)
         self._window: float = 0.0
@@ -188,19 +202,19 @@ def __repr__(self) -> str:
         return f"<Cooldown rate: {self.rate} per: {self.per} window: {self._window} tokens: {self._tokens}>"
 
 
-class CooldownMapping:
+class CooldownMapping(Generic[ContextTypeT]):
     def __init__(
         self,
         original: Cooldown | None,
-        type: Callable[[Message], Any],
+        type: Callable[[ContextTypeT], Any],
     ) -> None:
         if not callable(type):
             msg = "Cooldown type must be a BucketType or callable"
             raise TypeError(msg)
 
         self._cache: dict[Any, Cooldown] = {}
         self._cooldown: Cooldown | None = original
-        self._type: Callable[[Message], Any] = type
+        self._type: Callable[[ContextTypeT], Any] = type
 
     def copy(self) -> CooldownMapping:
         ret = CooldownMapping(self._cooldown, self._type)
@@ -212,14 +226,14 @@ def valid(self) -> bool:
         return self._cooldown is not None
 
     @property
-    def type(self) -> Callable[[Message], Any]:
+    def type(self) -> Callable[[ContextTypeT], Any]:
         return self._type
 
     @classmethod
-    def from_cooldown(cls, rate: float, per: float, type) -> Self:
+    def from_cooldown(cls, rate: int, per: float, type) -> Self:
         return cls(Cooldown(rate, per), type)
 
-    def _bucket_key(self, msg: Message) -> Any:
+    def _bucket_key(self, msg: ContextTypeT) -> Any:
         return self._type(msg)
 
     def _verify_cache_integrity(self, current: float | None = None) -> None:
@@ -235,11 +249,11 @@ def _is_default(self) -> bool:
         # This method can be overridden in subclasses
         return self._type is BucketType.default
 
-    def create_bucket(self, message: Message) -> Cooldown:
+    def create_bucket(self, message: ContextTypeT) -> Cooldown:
         assert self._cooldown is not None
         return self._cooldown.copy()
 
-    def get_bucket(self, message: Message, current: float | None = None) -> Cooldown:
+    def get_bucket(self, message: ContextTypeT, current: float | None = None) -> Cooldown:
         if self._is_default():
             assert self._cooldown is not None
             return self._cooldown
@@ -255,17 +269,19 @@ def get_bucket(self, message: Message, current: float | None = None) -> Cooldown
 
         return bucket
 
-    def update_rate_limit(self, message: Message, current: float | None = None) -> float | None:
+    def update_rate_limit(
+        self, message: ContextTypeT, current: float | None = None
+    ) -> float | None:
         bucket = self.get_bucket(message, current)
         return bucket.update_rate_limit(current)
 
 
-class DynamicCooldownMapping(CooldownMapping):
+class DynamicCooldownMapping(CooldownMapping[ContextTypeT]):
     def __init__(
-        self, factory: Callable[[Message], Cooldown], type: Callable[[Message], Any]
+        self, factory: Callable[[ContextTypeT], Cooldown], type: Callable[[ContextTypeT], Any]
     ) -> None:
         super().__init__(None, type)
-        self._factory: Callable[[Message], Cooldown] = factory
+        self._factory: Callable[[ContextTypeT], Cooldown] = factory
 
     def copy(self) -> DynamicCooldownMapping:
         ret = DynamicCooldownMapping(self._factory, self._type)
@@ -280,7 +296,7 @@ def _is_default(self) -> bool:
         # In dynamic mappings even default bucket types may have custom behavior
         return False
 
-    def create_bucket(self, message: Message) -> Cooldown:
+    def create_bucket(self, message: ContextTypeT) -> Cooldown:
         return self._factory(message)
 
 
@@ -302,7 +318,7 @@ class _Semaphore:
     def __init__(self, number: int) -> None:
         self.value: int = number
         self.loop: asyncio.AbstractEventLoop = asyncio.get_running_loop()
-        self._waiters: deque[asyncio.Future] = deque()
+        self._waiters: deque[asyncio.Future[None]] = deque()
 
     def __repr__(self) -> str:
         return f"<_Semaphore value={self.value} waiters={len(self._waiters)}>"
@@ -367,10 +383,10 @@ def copy(self) -> Self:
     def __repr__(self) -> str:
         return f"<MaxConcurrency per={self.per!r} number={self.number} wait={self.wait}>"
 
-    def get_key(self, message: Message) -> Any:
+    def get_key(self, message: ContextType) -> Any:
         return self.per.get_key(message)
 
-    async def acquire(self, message: Message) -> None:
+    async def acquire(self, message: ContextType) -> None:
         key = self.get_key(message)
 
         try:
@@ -382,7 +398,7 @@ async def acquire(self, message: Message) -> None:
         if not acquired:
             raise MaxConcurrencyReached(self.number, self.per)
 
-    async def release(self, message: Message) -> None:
+    async def release(self, message: ContextType) -> None:
         # Technically there's no reason for this function to be async
         # But it might be more useful in the future
         key = self.get_key(message)