allow to download multiple media simultaneously

Author: RuslanUC

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "t-export"
-version = "0.1.3b3"
+version = "0.1.3b4"
 description = "Telegram chats export tool."
 authors = ["RuslanUC <dev_ruslan_uc@protonmail.com>"]
 readme = "README.md"

texport/export_config.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
 from typing import Union
@@ -32,10 +32,15 @@ class ExportConfig:
     to_date: datetime = datetime.now()
     print: bool = False
     preload: bool = False
+    max_concurrent_downloads: int = 4
 
     def excluded_media(self) -> set[MessageMediaType]:
         result = set()
         for media_type in EXPORT_MEDIA:
             if not getattr(self, f"export_{media_type}"):
                 result.add(EXPORT_MEDIA[media_type])
         return result
+
+    def __post_init__(self):
+        if self.max_concurrent_downloads <= 0:
+            self.max_concurrent_downloads = 4

texport/main.py

@@ -52,10 +52,12 @@ async def _main(session_name: str, api_id: int, api_hash: str, config: ExportCon
 @click.option("--documents/--no-documents", default=True, help="Download documents or not.")
 @click.option("--quiet", "-q", is_flag=True, default=False, help="Do not print progress to console.")
 @click.option("--no-preload", is_flag=True, default=False, help="Do not preload all messages.")
+@click.option("--max-concurrent-downloads", "-d", type=click.INT, default=4,
+              help="Number of concurrent media downloads.")
 def main(
         session_name: str, api_id: int, api_hash: str, chat_id: str, output: str, size_limit: int, from_date: str,
         to_date: str, photos: bool, videos: bool, voice: bool, video_notes: bool, stickers: bool, gifs: bool,
-        documents: bool, quiet: bool, no_preload: bool,
+        documents: bool, quiet: bool, no_preload: bool, max_concurrent_downloads: int,
 ) -> None:
     home = Path.home()
     texport_dir = home / ".texport"
@@ -77,6 +79,7 @@ def main(
         export_files=documents,
         print=not quiet,
         preload=not no_preload,
+        max_concurrent_downloads=max_concurrent_downloads,
     )
 
     if session_name.endswith(".session"):

texport/media_downloader.py

@@ -1,5 +1,5 @@
 import asyncio
-from asyncio import sleep
+from asyncio import sleep, Task, Semaphore, create_task
 from os.path import relpath
 from typing import Union, Optional
 
@@ -22,6 +22,8 @@ def __init__(self, client: Client, config: ExportConfig, media_dict: dict, progr
         self.progress = progress
 
         self._running = False
+        self._downloading: dict[Union[str, int], ...] = {}
+        self._sem = Semaphore(self.config.max_concurrent_downloads)
 
     def add(self, file_id: str, download_dir: str, out_id: Union[str, int]) -> None:
         if out_id in self.all_ids: return
@@ -31,28 +33,34 @@ def add(self, file_id: str, download_dir: str, out_id: Union[str, int]) -> None:
         self._status()
 
     async def _download(self, file_id: str, download_dir: str, out_id: Union[str, int]) -> None:
-        try:
-            path = await self.client.download_media(file_id, file_name=download_dir)
-        except RPCError:
-            return
+        async with self._sem:
+            try:
+                path = await self.client.download_media(file_id, file_name=download_dir)
+            except RPCError:
+                return
+            finally:
+                self._downloading.pop(out_id, None)
+                self.ids.discard(out_id)
+
         path = relpath(path, self.config.output_dir.absolute())
         self.output[out_id] = path
 
     def _status(self, status: str=None) -> None:
         with self.progress.update():
             self.progress.media_status = status or self.progress.media_status
-            self.progress.media_queue = len(self.queue)
+            self.progress.media_queue = len(self.queue) + len(self._downloading)
 
     async def _task(self) -> None:
+        # use create_task and semaphore
+        downloading: dict[Union[str, int], Task] = {}
         while self._running:
-            if not self.queue:
+            if not self.queue and not downloading:
                 self._status("Idle...")
                 await sleep(.1)
                 continue
             self._status("Downloading...")
-            await self._download(*self.queue[0])
-            _, _, task_id = self.queue.pop(0)
-            self.ids.discard(task_id)
+            *args, task_id = self.queue.pop(0)
+            self._downloading[task_id] = create_task(self._download(*args, task_id))
 
         self._status("Stopped...")
 
@@ -66,7 +74,7 @@ async def stop(self) -> None:
 
     async def wait(self, messages: Optional[list[int]]=None) -> None:
         messages = set(messages) if messages is not None else None
-        while self._running and self.queue:
+        while self._running and (self.queue or self._downloading):
             if messages is not None and not messages.intersection(self.ids):
                 break
             await sleep(.1)