Add command cooldowns implementation

Author: EvieePy

twitchio/ext/commands/__init__.py

@@ -25,5 +25,6 @@
 from .bot import Bot as Bot
 from .components import *
 from .context import *
+from .cooldowns import *
 from .core import *
 from .exceptions import *

twitchio/ext/commands/cooldowns.py

@@ -0,0 +1,280 @@
+"""
+MIT License
+
+Copyright (c) 2017 - Present PythonistaGuild
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+from __future__ import annotations
+
+import abc
+import asyncio
+import datetime
+import enum
+from collections.abc import Callable, Coroutine, Hashable
+from typing import TYPE_CHECKING, Any, Generic, Self, TypeAlias, TypeVar
+
+
+if TYPE_CHECKING:
+    import twitchio
+
+    from .context import Context
+
+
+__all__ = ("BaseCooldown", "Cooldown", "GCRACooldown", "BucketType", "Bucket")
+
+
+PT = TypeVar("PT")
+CT = TypeVar("CT")
+
+
+class BucketType(enum.Enum):
+    """Enum representing default implementations for the key argument in :func:`~.commands.cooldown`.
+
+    Attributes
+    ----------
+    default
+        The cooldown will be considered a global cooldown shared across every channel and user.
+    user
+        The cooldown will apply per user, accross all channels.
+    channel
+        The cooldown will apply to every user/chatter in the channel.
+    chatter
+        The cooldown will apply per user, per channel.
+    """
+
+    default = 0
+    user = 1
+    channel = 2
+    chatter = 3
+
+    def get_key(self, payload: twitchio.ChatMessage | Context) -> Any:
+        if self is BucketType.user:
+            return payload.chatter.id
+
+        elif self is BucketType.channel:
+            return ("channel", payload.broadcaster.id)
+
+        elif self is BucketType.chatter:
+            return (payload.broadcaster.id, payload.chatter.id)
+
+    def __call__(self, payload: twitchio.ChatMessage | Context) -> Any:
+        return self.get_key(payload)
+
+
+class BaseCooldown(abc.ABC):
+    """Base class used to implement your own cooldown algorithm for use with :func:`~.commands.cooldown`.
+
+    Some built-in cooldown algorithms already exist:
+
+    - :class:`~.commands.Cooldown` - (``Token Bucket Algorithm``)
+
+    - :class:`~.commands.GCRACooldown` - (``Generic Cell Rate Algorithm``)
+
+
+    .. note::
+
+        Every base method must be implemented in this base class.
+    """
+
+    @abc.abstractmethod
+    def reset(self) -> None:
+        """Base method which should be implemented to reset the cooldown."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def update(self) -> float | None:
+        """Base method which should be implemented to update the cooldown/ratelimit.
+
+        This is where your algorithm logic should be contained.
+
+        .. important::
+
+            This method should always return a :class:`float` or ``None``. If ``None`` is returned by this method,
+            the cooldown will be considered bypassed.
+
+        Returns
+        -------
+        :class:`float`
+            The time needed to wait before you are off cooldown.
+        ``None``
+            Bypasses the cooldown.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def copy(self) -> Self:
+        """Base method which should be implemented to return a copy of this class in it's original state."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def is_ratelimited(self) -> bool:
+        """Base method which should be implemented which returns a bool indicating whether the cooldown is ratelimited.
+
+        Returns
+        -------
+        bool
+            A bool indicating whether this cooldown is currently ratelimited.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def is_dead(self) -> bool:
+        """Base method which should be implemented to indicate whether the cooldown should be considered stale and allowed
+        to be removed from the ``bucket: cooldown`` mapping.
+
+        Returns
+        -------
+        bool
+            A bool indicating whether this cooldown is stale/old.
+        """
+        raise NotImplementedError
+
+
+class Cooldown(BaseCooldown):
+    """Default cooldown algorithm for :func:`~.commands.cooldown`, which implements a ``Token Bucket Algorithm``.
+
+    See: :func:`~.commands.cooldown` for more documentation.
+    """
+
+    def __init__(self, *, rate: int, per: int | datetime.timedelta) -> None:
+        if rate <= 0:
+            raise ValueError(f'Cooldown rate must be equal to or greater than 1. Got "{rate}" expected >= 1.')
+
+        self._rate: int = rate
+        self._per: datetime.timedelta = datetime.timedelta(seconds=per) if isinstance(per, int) else per
+
+        now: datetime.datetime = datetime.datetime.now(tz=datetime.UTC)
+        self._window: datetime.datetime = now + self._per
+
+        if self._window <= now:
+            raise ValueError("The provided per value for Cooldowns can not go into the past.")
+
+        self._tokens: int = self._rate
+        self.last_updated: datetime.datetime | None = None
+
+    @property
+    def per(self) -> datetime.timedelta:
+        return self._per
+
+    def reset(self) -> None:
+        self._tokens = self._rate
+        self._window = datetime.datetime.now(tz=datetime.UTC) + self._per
+
+    def get_tokens(self, now: datetime.datetime | None = None) -> int:
+        if now is None:
+            now = datetime.datetime.now(tz=datetime.UTC)
+
+        tokens = max(self._tokens, 0)
+        if now > self._window:
+            tokens = self._rate
+
+        return tokens
+
+    def is_ratelimited(self) -> bool:
+        self._tokens = self.get_tokens()
+        return self._tokens == 0
+
+    def update(self, *, factor: int = 1) -> float | None:
+        now = datetime.datetime.now(tz=datetime.UTC)
+        self.last_updated = now
+
+        self._tokens = self.get_tokens(now)
+
+        if self._tokens == self._rate:
+            self._window = datetime.datetime.now(tz=datetime.UTC) + self._per
+
+        self._tokens -= factor
+
+        if self._tokens < 0:
+            remaining = (self._window - now).total_seconds()
+            return remaining
+
+    def copy(self) -> Self:
+        return self.__class__(rate=self._rate, per=self._per)
+
+    def is_dead(self) -> bool:
+        if self.last_updated is None:
+            return False
+
+        now = datetime.datetime.now(tz=datetime.UTC)
+        return now > (self.last_updated + self.per)
+
+
+class GCRACooldown(BaseCooldown):
+    """Not implemented yet."""
+
+
+KeyT: TypeAlias = Callable[[Any], Hashable] | Callable[[Any], Coroutine[Any, Any, Hashable]] | BucketType
+
+
+class Bucket(Generic[PT]):
+    def __init__(self, cooldown: BaseCooldown, *, key: KeyT) -> None:
+        self._cooldown: BaseCooldown = cooldown
+        self._cache: dict[Hashable, BaseCooldown] = {}
+        self._key: KeyT = key
+
+    @classmethod
+    def from_cooldown(cls, *, base: type[BaseCooldown], key: KeyT, **kwargs: Any) -> Self:
+        cd: BaseCooldown = base(**kwargs)
+        return cls(cd, key=key)
+
+    def create_cooldown(self) -> BaseCooldown | None:
+        return self._cooldown.copy()
+
+    def verify_cache(self) -> None:
+        dead = [k for k, v in self._cache.items() if v.is_dead()]
+        for key in dead:
+            del self._cache[key]
+
+    async def get_key(self, payload: PT) -> Hashable:
+        if asyncio.iscoroutinefunction(self._key):
+            key = await self._key(payload)
+        else:
+            key = self._key(payload)  # type: ignore
+
+        return key
+
+    async def get_cooldown(self, payload: PT) -> BaseCooldown | None:
+        if self._key is BucketType.default:
+            return self._cooldown
+
+        self.verify_cache()
+        key = await self.get_key(payload)
+        if key is None:
+            return
+
+        if key not in self._cache:
+            cooldown = self.create_cooldown()
+
+            if cooldown is not None:
+                self._cache[key] = cooldown
+        else:
+            cooldown = self._cache[key]
+
+        return cooldown
+
+    async def update(self, payload: PT, **kwargs: Any) -> float | None:
+        bucket = await self.get_cooldown(payload)
+
+        if bucket is None:
+            return None
+
+        return bucket.update(**kwargs)