changes to leaky bucket

datawhores · Sep 11, 2024 · 25e5a2d · 25e5a2d
1 parent 42c4438
commit 25e5a2d
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/ofscraper/classes/sessionmanager/download.py b/ofscraper/classes/sessionmanager/download.py
@@ -103,13 +103,14 @@ async def factoryasync(self, input):
     def chunk_with_limit(self, funct):
         async def wrapper(*args, **kwargs):
             async for chunk in funct(*args, **kwargs):
-                size = len(chunk)
-                await self.get_token(size)
+                await self.get_token(chunk)
                 yield chunk
         return wrapper
 
-    async def get_token(self, size):
-        await self.leaky_bucket.acquire(size)
+    async def get_token(self, chunk):
+        if settings.get_download_limit() <= 0:
+            return
+        await self.leaky_bucket.acquire(chunk)
 
 
 class cdm_session(sessionManager.sessionManager):