simpfly speed limit logic and edge cases

Author: mjishnu

pypdl/consumer.py

@@ -3,7 +3,7 @@
 
 from aiofiles import os
 
-from .downloader import Multidown, Singledown
+from .downloader import SegmentDownloader, SingleSegmentDownloader
 from .utils import (
     FileValidator,
     auto_cancel_gather,
@@ -120,7 +120,7 @@ async def _multi_segment(self, segment_table, file_path, speed_limit, **kwargs):
         speed_limit = speed_limit / segments
         self._logger.debug("Multi-Segment download started %s", self._id)
         for segment in range(segments):
-            md = Multidown(self._session, speed_limit)
+            md = SegmentDownloader(self._session, speed_limit)
             self._workers.append(md)
             tasks.add(asyncio.create_task(md.worker(segment_table, segment, **kwargs)))
 
@@ -132,7 +132,7 @@ async def _multi_segment(self, segment_table, file_path, speed_limit, **kwargs):
 
     async def _single_segment(self, url, file_path, speed_limit, **kwargs):
         self._logger.debug("Single-Segment download started %s", self._id)
-        sd = Singledown(self._session, speed_limit)
+        sd = SingleSegmentDownloader(self._session, speed_limit)
         self._workers.append(sd)
         await sd.worker(url, file_path, **kwargs)
         self._logger.debug("Downloaded single segment %s", self._id)

pypdl/downloader.py

@@ -7,44 +7,39 @@
 MEGABYTE = 1048576
 
 
-class Basicdown:
+class BaseDownloader:
     """Base downloader class."""
 
     def __init__(self, session: ClientSession, speed_limit: float) -> None:
         self.session = session
-        self.speed_limit = speed_limit * MEGABYTE
+        self.speed_limit = max(0, speed_limit * MEGABYTE)
         self.curr = 0
 
     async def download(self, url: str, path: str, mode: str, **kwargs) -> None:
         """Download data in chunks."""
-        speedlimit_time = time.time()
-        speedlimit_size = 0
+        start_time = time.monotonic()
         async with self.session.get(url, **kwargs) as response:
             async with aiofiles.open(path, mode) as file:
                 async for chunk in response.content.iter_chunked(MEGABYTE):
-                    if self.speed_limit > 0:
-                        now = time.time()
-                        time_passed = now - speedlimit_time
-                        if time_passed > 0.1:
-                            curr_download = self.curr - speedlimit_size
-                            if curr_download / time_passed >= self.speed_limit:
-                                await asyncio.sleep(curr_download / self.speed_limit)
-                            else:
-                                speedlimit_time = now
-                                speedlimit_size = self.curr
-
                     await file.write(chunk)
                     self.curr += len(chunk)
 
+                    if self.speed_limit > 0:
+                        expected_time = self.curr / self.speed_limit
+                        current_time = time.monotonic() - start_time
+                        sleep_time = expected_time - current_time
+                        if sleep_time > 0:
+                            await asyncio.sleep(sleep_time)
+
 
-class Singledown(Basicdown):
+class SingleSegmentDownloader(BaseDownloader):
     """Class for downloading the whole file in a single segment."""
 
     async def worker(self, url: str, file_path: str, **kwargs) -> None:
         await self.download(url, file_path, "wb", **kwargs)
 
 
-class Multidown(Basicdown):
+class SegmentDownloader(BaseDownloader):
     """Class for downloading a specific segment of the file."""
 
     async def worker(self, segment_table: dict, id: int, **kwargs) -> None: