add ability to export multiple chats simultaneously

Author: RuslanUC

.gitignore

@@ -1,3 +1,4 @@
 .idea
 __pycache__
-dist
+dist
+/telegram_export

pyproject.toml

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

texport/__init__.py

@@ -0,0 +1,6 @@
+from .export_config import ExportConfig
+from .progress_print import ProgressPrint
+from .media_downloader import MediaExporter
+from .messages_preloader import Preloader
+from .messages_saver import MessagesSaver
+from .exporter import Exporter

texport/export_config.py

@@ -18,7 +18,7 @@
 
 @dataclass
 class ExportConfig:
-    chat_id: Union[str, int] = "me"
+    chat_ids: list[Union[str, int]] = field(default_factory=lambda: ["me"])
     output_dir: Path = Path("./telegram_export")
     export_photos: bool = True
     export_videos: bool = True

texport/exporter.py

@@ -6,12 +6,8 @@
 from pyrogram.types import Message as PyroMessage
 from pyrogram.utils import zero_datetime
 
-from .export_config import ExportConfig
+from . import ExportConfig, MediaExporter, Preloader, MessagesSaver, ProgressPrint
 from .media import MEDIA_TYPES
-from .media_downloader import MediaExporter
-from .messages_preloader import Preloader
-from .messages_saver import MessagesSaver
-from .progress_print import ProgressPrint
 
 
 class Exporter:
@@ -35,50 +31,54 @@ async def _export_media(self, message: PyroMessage) -> None:
             return
 
         if m.downloadable:
-            self._media_downloader.add(media.file_id, f"{self._config.output_dir.absolute()}/{m.dir_name}/", message.id)
+            chat_output_dir = (self._config.output_dir / f"{message.chat.id}").absolute()
 
+            self._media_downloader.add(media.file_id, f"{chat_output_dir}/{m.dir_name}/", message.id)
             if hasattr(media, "thumbs") and media.thumbs:
-                self._media_downloader.add(media.thumbs[0].file_id, f"{self._config.output_dir.absolute()}/thumbs/",
-                                           f"{message.id}_thumb")
+                self._media_downloader.add(media.thumbs[0].file_id, f"{chat_output_dir}/thumbs/", f"{message.id}_thumb")
 
     async def _write(self, wait_media: list[int]) -> None:
         self.progress.status = "Waiting for all media to be downloaded..."
         await self._media_downloader.wait(wait_media)
         self.progress.status = "Writing messages to file..."
         await self._saver.save()
 
-    async def _export(self, chat_id: Union[int, str]):
+    async def _export(self):
         await self._media_downloader.run()
 
         offset_date = zero_datetime() if self._config.to_date.date() >= date.today() else self._config.to_date
         loaded = 0
         medias = []
-        self.progress.approx_messages_count = await self._client.get_chat_history_count(chat_id)
-        messages_iter = Preloader(self._client, self.progress, self._export_media) \
+        for chat_id in self._config.chat_ids:
+            self.progress.approx_messages_count += await self._client.get_chat_history_count(chat_id)
+        messages_iter = Preloader(self._client, self.progress, self._config.chat_ids, self._export_media) \
             if self._config.preload else self._client.get_chat_history
-        async for message in messages_iter(chat_id, offset_date=offset_date):
-            if message.date < self._config.from_date:
-                break
 
-            loaded += 1
-            with self.progress.update():
-                self.progress.status = "Exporting messages..."
-                self.progress.messages_exported = loaded
+        for chat_id in self._config.chat_ids:
+            async for message in messages_iter(chat_id, offset_date=offset_date):
+                if message.date < self._config.from_date:
+                    break
 
-            if message.media:
-                medias.append(message.id)
-                medias.append(f"{message.id}_thumb")
-                await self._export_media(message)
+                loaded += 1
+                with self.progress.update():
+                    self.progress.status = "Exporting messages..."
+                    self.progress.messages_exported = loaded
 
-            if not message.text and not message.caption and message.media not in MEDIA_TYPES:
-                continue
+                if message.media:
+                    medias.append(message.id)
+                    medias.append(f"{message.id}_thumb")
+                    await self._export_media(message)
 
-            self._messages.append(message)
-            if len(self._messages) > 1000:
+                if not message.text and not message.caption and message.media not in MEDIA_TYPES:
+                    continue
+
+                self._messages.append(message)
+                if len(self._messages) > 1000:
+                    await self._write(medias)
+
+            if self._messages:
                 await self._write(medias)
 
-        if self._messages:
-            await self._write(medias)
         self._task = None
 
         self.progress.status = "Stopping media downloader..."
@@ -88,7 +88,7 @@ async def _export(self, chat_id: Union[int, str]):
     async def export(self, block: bool=True) -> None:
         if self._task is not None:
             return
-        coro = self._export(self._config.chat_id)
+        coro = self._export()
         if block:
             await coro
         else:

texport/main.py

@@ -33,7 +33,7 @@ async def _main(session_name: str, api_id: int, api_hash: str, config: ExportCon
               help="Telegram api hash. Saved in ~/.texport/config.json file.")
 @click.option("--session-name", "-s", type=click.STRING, default="main",
               help="Pyrogram session name or path to existing file. Saved in ~/.texport/<session_name>.session file.")
-@click.option("--chat-id", "-c", type=click.STRING, default="me",
+@click.option("--chat-id", "-c", type=click.STRING, default=["me"], multiple=True,
               help="Chat id or username or phone number. \"me\" or \"self\" to export saved messages.")
 @click.option("--output", "-o", type=click.STRING, default="./telegram_export",
               help="Output directory.")
@@ -55,7 +55,7 @@ async def _main(session_name: str, api_id: int, api_hash: str, config: ExportCon
 @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,
+        session_name: str, api_id: int, api_hash: str, chat_id: list[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, max_concurrent_downloads: int,
 ) -> None:
@@ -65,7 +65,7 @@ def main(
     makedirs(output, exist_ok=True)
 
     config = ExportConfig(
-        chat_id=chat_id,
+        chat_ids=chat_id,
         output_dir=Path(output),
         size_limit=size_limit,
         from_date=datetime.strptime(from_date, "%d.%m.%Y"),

texport/media_downloader.py

@@ -1,6 +1,7 @@
 import asyncio
 from asyncio import sleep, Task, Semaphore, create_task
 from os.path import relpath
+from pathlib import Path
 from typing import Union, Optional
 
 from pyrogram import Client
@@ -42,8 +43,7 @@ async def _download(self, file_id: str, download_dir: str, out_id: Union[str, in
                 self._downloading.pop(out_id, None)
                 self.ids.discard(out_id)
 
-        path = relpath(path, self.config.output_dir.absolute())
-        self.output[out_id] = path
+        self.output[out_id] = relpath(path, Path(download_dir).parent.absolute())
 
     def _status(self, status: str=None) -> None:
         with self.progress.update():

texport/messages_preloader.py

@@ -1,25 +1,28 @@
 from asyncio import sleep, get_event_loop
 
 from pyrogram import Client
-from pyrogram.types import Message as PyroMessage
+from pyrogram.types import Message
 
 from .progress_print import ProgressPrint
 
 
 class Preloader:
-    def __init__(self, client: Client, progress: ProgressPrint, media_cb):
+    def __init__(self, client: Client, progress: ProgressPrint, chat_ids: list, media_cb):
         self.client = client
         self.progress = progress
-        self.finished = False
-        self.messages: list[PyroMessage] = []
+        self.finished = {chat_id: False for chat_id in chat_ids}
+        self.messages: dict[..., list[Message]] = {chat_id: [] for chat_id in chat_ids}
         self.messages_loaded = 0
         self.media_cb = media_cb
+        self._chat_ids = chat_ids
 
         self._task = None
         self._pyro_args = ()
         self._pyro_kwargs = {}
+        self._current_chat_id = None
 
-    def __call__(self, *pyrogram_args, **pyrogram_kwargs):
+    def __call__(self, chat_id, *pyrogram_args, **pyrogram_kwargs):
+        self._current_chat_id = chat_id
         self._pyro_args = pyrogram_args
         self._pyro_kwargs = pyrogram_kwargs
         return self
@@ -28,27 +31,28 @@ def __aiter__(self):
         return self
 
     async def _preload(self) -> None:
-        async for message in self.client.get_chat_history(*self._pyro_args, **self._pyro_kwargs):
-            self.messages.append(message)
-            self.messages_loaded += 1
+        for chat_id in self._chat_ids:
+            async for message in self.client.get_chat_history(chat_id, *self._pyro_args, **self._pyro_kwargs):
+                self.messages[chat_id].append(message)
+                self.messages_loaded += 1
 
-            if message.media and self.media_cb:
-                await self.media_cb(message)
+                if message.media and self.media_cb:
+                    await self.media_cb(message)
 
-            with self.progress.update():
-                self.progress.status = "Preloading messages and media..."
-                self.progress.messages_loaded = self.messages_loaded
+                with self.progress.update():
+                    self.progress.status = "Preloading messages and media..."
+                    self.progress.messages_loaded = self.messages_loaded
 
-        self.finished = True
+            self.finished[chat_id] = True
 
-    async def __anext__(self) -> PyroMessage:
+    async def __anext__(self) -> Message:
         if self._task is None: self._task = get_event_loop().create_task(self._preload())
 
-        while not self.finished and not self.messages:
+        while not self.finished[self._current_chat_id] and not self.messages[self._current_chat_id]:
             await sleep(.01)
 
-        if self.finished and not self.messages:
+        if self.finished[self._current_chat_id] and not self.messages[self._current_chat_id]:
             raise StopAsyncIteration
 
-        return self.messages.pop(0)
+        return self.messages[self._current_chat_id].pop(0)
 

texport/messages_saver.py

@@ -19,7 +19,10 @@ def __init__(self, messages: list[PyroMessage], media: dict[Union[int, str], str
 
     def _save(self) -> None:
         out_dir = self.config.output_dir
-        if not exists(out_dir / "js") or exists(out_dir / "images") or exists(out_dir / "css"):
+        if self.messages:
+            out_dir = out_dir / str(self.messages[0].chat.id)
+
+        if not exists(out_dir / "js") or not exists(out_dir / "images") or not exists(out_dir / "css"):
             unpack_to(out_dir)
 
         output = ""