Added github integration

buffteks/lib/python3.11/site-packages/limiter/__init__.py

@@ -0,0 +1,6 @@
+from typing import TypeAlias
+
+from .limiter import *
+
+
+limit: TypeAlias = Limiter  # alias to create anonymous or single-use `Limiter`s

buffteks/lib/python3.11/site-packages/limiter/base.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+from typing import Final, TypeVar, ParamSpec, Callable, Awaitable
+from random import random, randrange
+
+from token_bucket import Limiter as TokenBucket, MemoryStorage
+
+
+WAKE_UP: Final[int] = 0
+
+
+P = ParamSpec('P')
+T = TypeVar('T')
+
+Decoratable = Callable[P, T] | Callable[P, Awaitable[T]]
+Decorated = Decoratable
+Decorator = Callable[[Decoratable[P, T]], Decorated[P, T]]
+
+Bucket = bytes
+BucketName = Bucket | str
+Num = int | float
+Tokens = Num
+Duration = Num
+UnitsInSecond = Duration
+JitterRange = range | tuple[int, int] | tuple[int, int, int]
+Jitter = int | bool | JitterRange
+
+
+CONSUME_TOKENS: Final[Tokens] = 1
+RATE: Final[Tokens] = 2
+CAPACITY: Final[Tokens] = 3
+MS_IN_SEC: Final[UnitsInSecond] = 1000
+
+DEFAULT_BUCKET: Final[Bucket] = b"default"
+DEFAULT_JITTER: Final[Jitter] = False
+
+
+def _get_bucket(name: BucketName) -> Bucket:
+  match name:
+    case bytes():
+      return name
+
+    case str():
+      return name.encode()
+
+  raise TypeError('Name must be bytes or a bytes-encodable string.')
+
+
+def _get_limiter(rate: Tokens = RATE, capacity: Tokens = CAPACITY) -> TokenBucket:
+  """
+  Returns TokenBucket object that implements a token-bucket algorithm.
+  """
+
+  return TokenBucket(rate, capacity, MemoryStorage())
+
+
+def _get_bucket_limiter(bucket: BucketName, limiter: 'Limiter') -> tuple[Bucket, TokenBucket]:
+  bucket: Bucket = _get_bucket(bucket)
+
+  if not isinstance(limiter, TokenBucket):
+    limiter: TokenBucket = limiter.limiter
+
+  return bucket, limiter
+
+
+def _get_sleep_duration(
+  consume: Tokens,
+  tokens: Tokens,
+  rate: Tokens,
+  jitter: Jitter = DEFAULT_JITTER,
+  units: UnitsInSecond = MS_IN_SEC
+) -> Duration:
+  """Increase contention by adding jitter to sleep duration"""
+  duration: Duration = (consume - tokens) / rate
+
+  match jitter:
+    case int() | float():
+      return duration - jitter
+
+    case bool() if jitter:
+      amount: Duration = random() / units
+      return duration - amount
+
+    case range():
+      amount: Duration = randrange(jitter.start, jitter.stop, jitter.step) / units
+      return duration - amount
+
+    case start, end:
+      amount: Duration = randrange(start, end) / units
+      return duration - amount
+
+    case start, end, step:
+      amount: Duration = randrange(start, end, step) / units
+      return duration - amount
+
+  return duration

buffteks/lib/python3.11/site-packages/limiter/limiter.py

@@ -0,0 +1,258 @@
+from __future__ import annotations
+
+import logging
+from typing import (
+  AsyncContextManager, ContextManager, Awaitable, TypedDict,
+  cast
+)
+from contextlib import (
+  AbstractContextManager, AbstractAsyncContextManager,
+  contextmanager, asynccontextmanager
+)
+from asyncio import sleep as aiosleep, iscoroutinefunction
+from dataclasses import dataclass, asdict
+from functools import wraps
+from enum import auto
+from time import sleep
+from abc import ABC
+import logging
+
+from strenum import StrEnum  # type: ignore
+from token_bucket import Limiter as TokenBucket  # type: ignore
+
+from .base import (
+  MS_IN_SEC, UnitsInSecond, WAKE_UP, RATE, CAPACITY, CONSUME_TOKENS, DEFAULT_BUCKET,
+  Tokens, Decoratable, Decorated, Decorator, Jitter, P, T,
+  BucketName, _get_limiter, _get_bucket_limiter,
+  _get_sleep_duration, DEFAULT_JITTER,
+)
+
+
+log = logging.getLogger(__name__)
+
+
+class Attrs(TypedDict):
+  consume: Tokens
+  bucket: BucketName
+
+  limiter: TokenBucket
+  jitter: Jitter
+  units: UnitsInSecond
+
+
+class LimiterBase(ABC):
+  consume: Tokens
+  bucket: BucketName
+
+  limiter: TokenBucket
+  jitter: Jitter
+  units: UnitsInSecond
+
+
+class LimiterContextManager(
+  LimiterBase,
+  AbstractContextManager,
+  AbstractAsyncContextManager
+):
+  def __enter__(self) -> ContextManager[Limiter]:
+    with limit_rate(self.limiter, self.consume, self.bucket, self.jitter, self.units) as limiter:
+      return limiter
+
+  def __exit__(self, *args):
+    pass
+
+  async def __aenter__(self) -> AsyncContextManager[Limiter]:
+    async with async_limit_rate(self.limiter, self.consume, self.bucket, self.jitter, self.units) as limiter:
+      return limiter
+
+  async def __aexit__(self, *args):
+    pass
+
+
+class AttrName(StrEnum):
+  rate: str = auto()
+  capacity: str = auto()
+
+  consume: str = auto()
+  bucket: str = auto()
+
+  limiter: str = auto()
+  jitter: str = auto()
+  units: str = auto()
+
+@dataclass
+class Limiter(LimiterContextManager):
+  rate: Tokens = RATE
+  capacity: Tokens = CAPACITY
+
+  consume: Tokens | None = None
+  bucket: BucketName = DEFAULT_BUCKET
+
+  limiter: TokenBucket | None = None
+  jitter: Jitter = DEFAULT_JITTER
+  units: UnitsInSecond = MS_IN_SEC
+
+  def __post_init__(self):
+    if self.limiter is None:
+      self.limiter = _get_limiter(self.rate, self.capacity)
+
+    if self.consume is None:
+      self.consume = CONSUME_TOKENS
+
+  def __call__(
+    self,
+    func_or_consume: Decoratable[P, T] | Tokens | None = None,
+    bucket: BucketName | None = None,
+    jitter: Jitter | None = None,
+    units: UnitsInSecond | None = None,
+    **attrs: Attrs,
+  ) -> Decorated[P, T] | Limiter:
+    if callable(func_or_consume):
+      func: Decoratable = cast(Decoratable, func_or_consume)
+      wrapper = limit_calls(self, self.consume, self.bucket)
+      return wrapper(func)
+
+    elif func_or_consume and not isinstance(func_or_consume, Tokens):
+      raise TypeError(f'First argument must be callable or {Tokens}')
+
+    if AttrName.rate in attrs or AttrName.capacity in attrs:
+      raise ValueError('Create a new limiter with the new() method or Limiter class')
+
+    consume: Tokens = cast(Tokens, func_or_consume)
+    new_attrs: Attrs = self.attrs
+
+    if consume:
+      new_attrs[AttrName.consume] = consume
+
+    if bucket:
+      new_attrs[AttrName.bucket] = bucket
+
+    if jitter:
+      new_attrs[AttrName.jitter] = jitter
+
+    if units:
+      new_attrs[AttrName.units] = units
+
+    new_attrs |= attrs
+
+    return Limiter(**new_attrs, limiter=self.limiter)
+
+  @property
+  def attrs(self) -> Attrs:
+    attrs = asdict(self)
+    attrs.pop(AttrName.limiter, None)
+
+    return attrs
+
+  def new(self, **attrs: Attrs):
+    new_attrs = self.attrs | attrs
+
+    return Limiter(**new_attrs)
+
+
+def limit_calls(
+  limiter: Limiter,
+  consume: Tokens = CONSUME_TOKENS,
+  bucket: BucketName = DEFAULT_BUCKET,
+  jitter: Jitter = DEFAULT_JITTER,
+  units: UnitsInSecond = MS_IN_SEC,
+) -> Decorator[P, T]:
+  """
+  Rate-limiting decorator for synchronous and asynchronous callables.
+  """
+  lim_wrapper: Limiter = limiter
+
+  bucket, limiter = _get_bucket_limiter(bucket, limiter)
+  limiter: TokenBucket = cast(TokenBucket, limiter)
+
+  def decorator(func: Decoratable[P, T]) -> Decorated[P, T]:
+    if iscoroutinefunction(func):
+      @wraps(func)
+      async def new_coroutine_func(*args: P.args, **kwargs: P.kwargs) -> Awaitable[T]:
+        async with async_limit_rate(limiter, consume, bucket, jitter, units):
+          return await func(*args, **kwargs)
+
+      new_coroutine_func.limiter = lim_wrapper
+      return new_coroutine_func
+
+    elif callable(func):
+      @wraps(func)
+      def new_func(*args: P.args, **kwargs: P.kwargs) -> T:
+        with limit_rate(limiter, consume, bucket, jitter, units):
+          return func(*args, **kwargs)
+
+      new_func.limiter = lim_wrapper
+      return new_func
+
+    else:
+      raise ValueError("Can only decorate callables and coroutine functions.")
+
+  return decorator
+
+
+@asynccontextmanager
+async def async_limit_rate(
+  limiter: Limiter,
+  consume: Tokens = CONSUME_TOKENS,
+  bucket: BucketName = DEFAULT_BUCKET,
+  jitter: Jitter = DEFAULT_JITTER,
+  units: UnitsInSecond = MS_IN_SEC,
+) -> AsyncContextManager[Limiter]:
+  """
+  Rate-limiting asynchronous context manager.
+  """
+  lim_wrapper: Limiter = limiter
+
+  bucket, limiter = _get_bucket_limiter(bucket, limiter)
+  limiter: TokenBucket = cast(TokenBucket, limiter)
+
+  # minimize attribute look ups in loop
+  get_tokens = limiter._storage.get_token_count
+  lim_consume = limiter.consume
+  rate = limiter._rate
+
+  while not lim_consume(bucket, consume):
+    tokens = get_tokens(bucket)
+    sleep_for = _get_sleep_duration(consume, tokens, rate, jitter, units)
+
+    if sleep_for <= WAKE_UP:
+      break
+
+    log.debug(f'Rate limit reached. Sleeping for {sleep_for}s.')
+    await aiosleep(sleep_for)
+
+  yield lim_wrapper
+
+
+@contextmanager
+def limit_rate(
+  limiter: Limiter,
+  consume: Tokens = CONSUME_TOKENS,
+  bucket: BucketName = DEFAULT_BUCKET,
+  jitter: Jitter = DEFAULT_JITTER,
+  units: UnitsInSecond = MS_IN_SEC,
+) -> ContextManager[Limiter]:
+  """
+  Thread-safe rate-limiting context manager.
+  """
+  lim_wrapper: Limiter = limiter
+
+  bucket, limiter = _get_bucket_limiter(bucket, limiter)
+  limiter: TokenBucket = cast(TokenBucket, limiter)
+
+  # minimize attribute look ups in loop
+  get_tokens = limiter._storage.get_token_count
+  lim_consume = limiter.consume
+  rate = limiter._rate
+
+  while not lim_consume(bucket, consume):
+    tokens = get_tokens(bucket)
+    sleep_for = _get_sleep_duration(consume, tokens, rate, jitter, units)
+
+    if sleep_for <= WAKE_UP:
+      break
+
+    log.debug(f'Rate limit reached. Sleeping for {sleep_for}s.')
+    sleep(sleep_for)
+
+  yield lim_wrapper