houdini/Houdini/Cooldown.py

import enum
import time


class CooldownError(Exception):
    """Raised when packets are sent whilst a cooldown is active"""
    pass


class BucketType(enum.Enum):
    Default = 1
    Penguin = 1
    Server = 2


class _Cooldown:

    __slots__ = ['rate', 'per', 'bucket_type', 'last',
                 '_window', '_tokens']

    def __init__(self, per, rate, bucket_type):
        self.per = float(per)
        self.rate = int(rate)
        self.bucket_type = bucket_type
        self.last = 0.0

        self._window = 0.0
        self._tokens = self.rate

    @property
    def is_cooling(self):
        current = time.time()
        self.last = current

        if self._tokens == self.rate:
            self._window = current

        if current > self._window + self.per:
            self._tokens = self.rate
            self._window = current

        if self._tokens == 0:
            return self.per - (current - self._window)

        self._tokens -= 1
        if self._tokens == 0:
            self._window = current

    def reset(self):
        self._tokens = self.rate
        self.last = 0.0

    def copy(self):
        return _Cooldown(self.per, self.rate, self.bucket_type)


class _CooldownMapping:

    __slots__ = ['_cooldown', '_cache', 'callback']

    def __init__(self, callback, cooldown_object):
        self._cooldown = cooldown_object

        self.callback = callback

        self._cache = {}

    def _get_bucket_key(self, p):
        if self._cooldown.bucket_type == BucketType.Default:
            return p
        return p.server

    def _verify_cache_integrity(self):
        current = time.time()
        self._cache = {cache_key: bucket for cache_key, bucket in
                       self._cache.items() if current < bucket.last + bucket.per}

    def get_bucket(self, p):
        self._verify_cache_integrity()
        cache_key = self._get_bucket_key(p)
        if cache_key not in self._cache:
            bucket = self._cooldown.copy()
            self._cache[cache_key] = bucket
        else:
            bucket = self._cache[cache_key]
        return bucket