v1.5

.github/workflows/test.yml

@@ -0,0 +1,27 @@
+name: Test Pypdl
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.8"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pytest
+        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    - name: Test with pytest
+      run: pytest test/test.py

pypdl/__init__.py

@@ -1,4 +1,3 @@
-__version__ = "1.4.5"
+__version__ = "1.5.0"
 
-from .pypdl_manager import Pypdl
-from .pypdl_factory import PypdlFactory
+from .pypdl import Pypdl

pypdl/consumer.py

@@ -0,0 +1,111 @@
+import asyncio
+
+from aiofiles import os
+
+from .downloader import Multidown, Singledown
+from .utils import FileValidator, combine_files, create_segment_table
+
+
+class Consumer:
+    def __init__(self, session, logger, _id):
+        self._workers = []
+        self._downloaded_size = 0
+        self._size = 0
+        self._show_size = True
+        self.id = _id
+
+        self.logger = logger
+        self.session = session
+        self.success = []
+
+    @property
+    def size(self):
+        if self._show_size:
+            self._size = (
+                sum(worker.curr for worker in self._workers) + self._downloaded_size
+            )
+        return self._size
+
+    async def process_tasks(self, in_queue, out_queue):
+        self.logger.debug("Consumer %s started", self.id)
+        while True:
+            task = await in_queue.get()
+            self.logger.debug("Consumer %s received task", self.id)
+            if task is None:
+                break
+            try:
+                await self._download(task)
+            except asyncio.CancelledError:
+                raise
+            except Exception as e:
+                self.logger.debug("Task %s failed", self.id)
+                self.logger.exception(e)
+                await out_queue.put([task[0]])
+
+            self._workers.clear()
+            self._show_size = True
+
+            self.logger.debug("Consumer %s completed task", self.id)
+        self.logger.debug("Consumer %s exited", self.id)
+        return self.success
+
+    async def _download(
+        self,
+        task,
+    ):
+        _id, task = task
+        (
+            url,
+            file_path,
+            multisegment,
+            etag,
+            size,
+            segments,
+            overwrite,
+            speed_limit,
+            etag_validation,
+            kwargs,
+        ) = task
+
+        self.logger.debug("Download started %s", self.id)
+        if not overwrite and await os.path.exists(file_path):
+            self.logger.debug("File already exists, download completed")
+            self.success.append(FileValidator(file_path))
+            self._downloaded_size += await os.path.getsize(file_path)
+            return
+
+        if multisegment:
+            segment_table = await create_segment_table(
+                url, file_path, segments, size, etag, etag_validation
+            )
+            await self._multi_segment(segment_table, file_path, speed_limit, **kwargs)
+        else:
+            await self._single_segment(url, file_path, speed_limit, **kwargs)
+
+        self.success.append((url, FileValidator(file_path)))
+        self.logger.debug("Download exited %s", self.id)
+
+    async def _multi_segment(self, segment_table, file_path, speed_limit, **kwargs):
+        tasks = set()
+        segments = segment_table["segments"]
+        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)
+            self._workers.append(md)
+            tasks.add(asyncio.create_task(md.worker(segment_table, segment, **kwargs)))
+
+        await asyncio.gather(*tasks)
+        await combine_files(file_path, segments)
+        self.logger.debug("Downloaded all segments %s", self.id)
+        self._show_size = False
+        self._downloaded_size += await os.path.getsize(file_path)
+
+    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)
+        self._workers.append(sd)
+        await sd.worker(url, file_path, **kwargs)
+        self.logger.debug("Downloaded single segment %s", self.id)
+        self._show_size = False
+        self._downloaded_size += await os.path.getsize(file_path)

pypdl/downloader.py

@@ -1,5 +1,5 @@
-from pathlib import Path
-from threading import Event
+import asyncio
+import time
 
 import aiofiles
 from aiohttp import ClientSession
@@ -10,68 +10,71 @@
 class Basicdown:
     """Base downloader class."""
 
-    def __init__(self, interrupt: Event):
-        self.curr = 0  # Downloaded size in bytes (current size)
+    def __init__(self, session: ClientSession, speed_limit: float) -> None:
+        self.session = session
+        self.speed_limit = speed_limit * MEGABYTE
+        self.curr = 0
         self.completed = False
-        self.interrupt = interrupt
-        self.downloaded = 0
 
-    async def download(
-        self, url: str, path: str, mode: str, session: ClientSession, **kwargs
-    ) -> None:
+    async def download(self, url: str, path: str, mode: str, **kwargs) -> None:
         """Download data in chunks."""
-        async with session.get(url, **kwargs) as response:
+        speedlimit_time = time.time()
+        speedlimit_size = 0
+        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)
-                    self.downloaded += len(chunk)
-                    if self.interrupt.is_set():
-                        break
 
 
 class Singledown(Basicdown):
     """Class for downloading the whole file in a single segment."""
 
-    async def worker(
-        self, url: str, file_path: str, session: ClientSession, **kwargs
-    ) -> None:
-        await self.download(url, file_path, "wb", session, **kwargs)
+    async def worker(self, url: str, file_path: str, **kwargs) -> None:
+        await self.download(url, file_path, "wb", **kwargs)
         self.completed = True
 
 
 class Multidown(Basicdown):
     """Class for downloading a specific segment of the file."""
 
-    async def worker(
-        self, segment_table: dict, id: int, session: ClientSession, **kwargs
-    ) -> None:
+    async def worker(self, segment_table: dict, id: int, **kwargs) -> None:
         url = segment_table["url"]
         overwrite = segment_table["overwrite"]
-        segment_path = Path(segment_table[id]["segment_path"])
-        start = segment_table[id]["start"]
-        end = segment_table[id]["end"]
+        segment_path = segment_table[id]["segment_path"]
         size = segment_table[id]["segment_size"]
 
-        if segment_path.exists():
-            downloaded_size = segment_path.stat().st_size
-            if overwrite or downloaded_size > size:
-                segment_path.unlink()
+        if await aiofiles.os.path.exists(segment_path):
+            downloaded_size = await aiofiles.os.path.getsize(segment_path)
+            if overwrite or downloaded_size > size.value:
+                await aiofiles.os.remove(segment_path)
             else:
                 self.curr = downloaded_size
 
         if kwargs.get("headers") is not None:
             kwargs["headers"] = kwargs["headers"].copy()
 
-        if self.curr < size:
-            start = start + self.curr
-            kwargs.setdefault("headers", {}).update({"range": f"bytes={start}-{end}"})
-            await self.download(url, segment_path, "ab", session, **kwargs)
+        if self.curr < size.value:
+            start = size.start + self.curr
+            kwargs.setdefault("headers", {}).update(
+                {"range": f"bytes={start}-{size.end}"}
+            )
+            await self.download(url, segment_path, "ab", **kwargs)
 
-        if self.curr == size:
+        if self.curr == size.value:
             self.completed = True
         else:
-            if not self.interrupt.is_set():
-                raise Exception(
-                    f"Incorrect segment size: expected {size} bytes, received {self.curr} bytes"
-                )
+            raise Exception(
+                f"Incorrect segment size: expected {size} bytes, received {self.curr} bytes"
+            )