diff --git a/disnake/ext/commands/base_core.py b/disnake/ext/commands/base_core.py index 6c146f54d9..9db8f2df43 100644 --- a/disnake/ext/commands/base_core.py +++ b/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) diff --git a/disnake/ext/commands/cooldowns.py b/disnake/ext/commands/cooldowns.py index afec770218..b11f930862 100644 --- a/disnake/ext/commands/cooldowns.py +++ b/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,11 +202,11 @@ def __repr__(self) -> str: return f"" -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" @@ -200,9 +214,9 @@ def __init__( 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: + def copy(self) -> CooldownMapping[ContextTypeT]: ret = CooldownMapping(self._cooldown, self._type) ret._cache = self._cache.copy() return ret @@ -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,19 +269,21 @@ 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: + def copy(self) -> DynamicCooldownMapping[ContextTypeT]: ret = DynamicCooldownMapping(self._factory, self._type) ret._cache = self._cache.copy() return ret @@ -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"" - 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) diff --git a/disnake/ext/commands/core.py b/disnake/ext/commands/core.py index b1baa351ac..c1fa7e134e 100644 --- a/disnake/ext/commands/core.py +++ b/disnake/ext/commands/core.py @@ -33,7 +33,14 @@ from .cog import Cog from .context import AnyContext, Context from .converter import Greedy, get_converter, run_converters -from .cooldowns import BucketType, Cooldown, CooldownMapping, DynamicCooldownMapping, MaxConcurrency +from .cooldowns import ( + BucketType, + ContextTypeT, + Cooldown, + CooldownMapping, + DynamicCooldownMapping, + MaxConcurrency, +) from .errors import ( ArgumentParsingError, BotMissingAnyRole, @@ -62,7 +69,7 @@ from typing_extensions import ParamSpec, Self, Unpack - from disnake.message import Message + from disnake import Message from ._types import AppCheck, Check, Coro, CoroFunc, Error, Hook @@ -77,7 +84,7 @@ class _CommandArgs(TypedDict, total=False): description: str hidden: bool checks: list[Check] - cooldown: CooldownMapping | None + cooldown: CooldownMapping[Message] | None max_concurrency: MaxConcurrency | None require_var_positional: bool ignore_extra: bool @@ -177,7 +184,7 @@ async def wrapped(*args: Any, **kwargs: Any) -> T | None: raise CommandInvokeError(exc) from exc finally: if command._max_concurrency is not None: - await command._max_concurrency.release(ctx) # pyright: ignore[reportArgumentType] + await command._max_concurrency.release(ctx) await command.call_after_hooks(ctx) return ret @@ -360,7 +367,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[Message] = buckets try: max_concurrency = func.__commands_max_concurrency__ @@ -801,7 +808,7 @@ def _prepare_cooldowns(self, ctx: Context) -> None: 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, ctx: Context) -> None: ctx.command = self @@ -811,8 +818,7 @@ async def prepare(self, ctx: Context) -> None: raise CheckFailure(msg) if self._max_concurrency is not None: - # For this application, context can be duck-typed as a Message - await self._max_concurrency.acquire(ctx) # pyright: ignore[reportArgumentType] + await self._max_concurrency.acquire(ctx) try: if self.cooldown_after_parsing: @@ -825,7 +831,7 @@ async def prepare(self, ctx: Context) -> None: await self.call_before_hooks(ctx) except Exception: if self._max_concurrency is not None: - await self._max_concurrency.release(ctx) # pyright: ignore[reportArgumentType] + await self._max_concurrency.release(ctx) raise def is_on_cooldown(self, ctx: Context) -> bool: @@ -2564,7 +2570,7 @@ def pred(ctx: AnyContext) -> bool: def cooldown( - rate: int, per: float, type: BucketType | Callable[[Message], Any] = BucketType.default + rate: int, per: float, type: BucketType | Callable[[ContextTypeT], Any] = BucketType.default ) -> Callable[[T], T]: r"""A decorator that adds a cooldown to a :class:`.Command` @@ -2585,7 +2591,7 @@ def cooldown( 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: :class:`.BucketType` | :class:`~collections.abc.Callable`\[[:class:`.Message`], :data:`~typing.Any`] + type: :class:`.BucketType` | :class:`~collections.abc.Callable`\[[:class:`.Message`], :data:`~typing.Any`] | :class:`~collections.abc.Callable`\[[:class:`.ApplicationCommandInteraction`], :data:`~typing.Any`] The type of cooldown to have. If callable, should return a key for the mapping. .. versionchanged:: 1.7 @@ -2596,7 +2602,8 @@ def decorator( func: Command[CogT, P, T] | CoroFunc, ) -> Command[CogT, P, T] | CoroFunc: if hasattr(func, "__command_flag__"): - func._buckets = CooldownMapping(Cooldown(rate, per), type) + # not assignable, because the decorator is typed Command[...] + func._buckets = CooldownMapping(Cooldown(rate, per), type) # pyright: ignore[reportAttributeAccessIssue] else: func.__commands_cooldown__ = CooldownMapping(Cooldown(rate, per), type) # pyright: ignore[reportAttributeAccessIssue] return func @@ -2605,7 +2612,7 @@ def decorator( def dynamic_cooldown( - cooldown: BucketType | Callable[[Message], Any], type: BucketType = BucketType.default + cooldown: BucketType | Callable[[ContextTypeT], Any], type: BucketType = BucketType.default ) -> Callable[[T], T]: r"""A decorator that adds a dynamic cooldown to a :class:`.Command` @@ -2643,7 +2650,8 @@ def decorator( func: Command[CogT, P, T] | CoroFunc, ) -> Command[CogT, P, T] | CoroFunc: if hasattr(func, "__command_flag__"): - func._buckets = DynamicCooldownMapping(cooldown, type) + # not assignable, because the decorator is typed Command[...] + func._buckets = DynamicCooldownMapping(cooldown, type) # pyright: ignore[reportAttributeAccessIssue] else: func.__commands_cooldown__ = DynamicCooldownMapping(cooldown, type) # pyright: ignore[reportAttributeAccessIssue] return func diff --git a/disnake/ext/commands/errors.py b/disnake/ext/commands/errors.py index 98252d3ebe..1a09e4bba7 100644 --- a/disnake/ext/commands/errors.py +++ b/disnake/ext/commands/errors.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from inspect import Parameter + from disnake import ApplicationCommandInteraction, Message from disnake.abc import GuildChannel from disnake.threads import Thread from disnake.types.snowflake import Snowflake, SnowflakeList @@ -648,10 +649,19 @@ class CommandOnCooldown(CommandError): The amount of seconds to wait before you can retry again. """ - def __init__(self, cooldown: Cooldown, retry_after: float, type: BucketType) -> None: + def __init__( + self, + cooldown: Cooldown, + retry_after: float, + type: BucketType + | Callable[[Message], Any] + | Callable[[ApplicationCommandInteraction], Any], + ) -> None: self.cooldown: Cooldown = cooldown self.retry_after: float = retry_after - self.type: BucketType = type + self.type: ( + BucketType | Callable[[Message], Any] | Callable[[ApplicationCommandInteraction], Any] + ) = type super().__init__(f"You are on cooldown. Try again in {retry_after:.2f}s") diff --git a/disnake/ext/commands/slash_core.py b/disnake/ext/commands/slash_core.py index e5de12f725..93a37acf75 100644 --- a/disnake/ext/commands/slash_core.py +++ b/disnake/ext/commands/slash_core.py @@ -386,7 +386,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) @@ -823,7 +823,7 @@ async def invoke(self, inter: ApplicationCommandInteraction) -> None: 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) diff --git a/disnake/utils.py b/disnake/utils.py index 5e82cba7b5..ed064044bb 100644 --- a/disnake/utils.py +++ b/disnake/utils.py @@ -134,7 +134,7 @@ def __get__(self, instance, owner): if TYPE_CHECKING: - from functools import cached_property as cached_property + cached_property: TypeAlias = property from .abc import Snowflake from .asset import AssetBytes