diff --git a/src/throttler.py b/src/throttler.py index 219e740..5ebfae0 100644 --- a/src/throttler.py +++ b/src/throttler.py @@ -1,33 +1,37 @@ import asyncio from asyncio_throttle import Throttler as AsyncThrottler + class GlobalThrottler: _instance = None _lock = asyncio.Lock() - + def __init__(self, rate_limit: int = 500, period: int = 60) -> None: - if not hasattr(self, '_initialized'): + if not hasattr(self, "_initialized"): self.rate_limit = rate_limit self.period = period - self.throttler = AsyncThrottler(rate_limit=self.rate_limit, period=self.period) + self.throttler = AsyncThrottler( + rate_limit=self.rate_limit, period=self.period + ) self._initialized = True - + @classmethod - async def get_instance(cls, rate_limit: int = 500, period: int = 60) -> 'GlobalThrottler': + async def get_instance( + cls, rate_limit: int = 500, period: int = 60 + ) -> "GlobalThrottler": async with cls._lock: if not cls._instance: cls._instance = cls(rate_limit, period) return cls._instance - + async def acquire(self): await self.throttler.acquire() - - def release(self): - self.throttler.release() - + async def __aenter__(self): await self.acquire() return self - + async def __aexit__(self, exc_type, exc_val, exc_tb): - self.release() + # asyncio_throttle.Throttler doesn't have a release() method + # It automatically manages rate limiting via acquire() + pass