add polls and contacts exporting

Author: RuslanUC

pyproject.toml

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

texport/exporter.py

@@ -3,6 +3,7 @@
 from typing import Union
 
 from pyrogram import Client
+from pyrogram.enums import MessageMediaType
 from pyrogram.types import Message as PyroMessage
 from pyrogram.utils import zero_datetime
 
@@ -31,14 +32,15 @@ async def _export_media(self, message: PyroMessage) -> None:
             return
         m = MEDIA_TYPES[message.media]
         media = m.get_media(message)
-        if media.file_size > self._config.size_limit * 1024 * 1024:
+        if m.has_size_limit and media.file_size > self._config.size_limit * 1024 * 1024:
             return
 
-        self._media_downloader.add(media.file_id, f"{self._config.output_dir.absolute()}/{m.dir_name}/", message.id)
+        if m.downloadable:
+            self._media_downloader.add(media.file_id, f"{self._config.output_dir.absolute()}/{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")
+            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")
 
     async def _write(self, wait_media: list[int]) -> None:
         self.progress.status = "Waiting for all media to be downloaded..."
@@ -69,7 +71,7 @@ async def _export(self, chat_id: Union[int, str]):
                 medias.append(f"{message.id}_thumb")
                 await self._export_media(message)
 
-            if not message.text and not message.caption and message.id not in self._media:
+            if not message.text and not message.caption and message.media not in MEDIA_TYPES:
                 continue
 
             self._messages.append(message)

texport/html/contact.py

@@ -0,0 +1,24 @@
+from pyrogram.types import Message as PyroMessage
+
+from .base import HtmlMedia
+
+
+class Contact(HtmlMedia):
+    def __init__(self, *args, message: PyroMessage):
+        self.contact = message.contact
+
+    def no_media(self) -> str:
+        return ""
+
+    def to_html(self) -> str:
+        last_name = "" if not self.contact.last_name else self.contact.last_name
+
+        return f"""
+        <div class="media clearfix pull_left media_contact">
+            <div class="fill pull_left"></div>
+            <div class="body">
+                <div class="title bold">{self.contact.first_name} {last_name}</div>
+                <div class="status details">{self.contact.phone_number}</div>
+            </div>
+        </div>
+        """

texport/html/poll.py

@@ -0,0 +1,57 @@
+from pyrogram.enums import PollType
+from pyrogram.types import Message as PyroMessage, Poll as PyroPoll
+
+from .base import HtmlMedia, BaseComponent
+
+
+class PollOption(BaseComponent):
+    def __init__(self, poll: PyroPoll, option: int):
+        self.poll = poll
+        self.option_id = option
+        self.option = poll.options[option]
+
+    def to_html(self) -> str:
+        details = []
+        if self.option.voter_count:
+            details.append(f"{self.option.voter_count} votes")
+        if self.option_id == self.poll.chosen_option_id:
+            details.append("chosen vote")
+        if self.poll.type == PollType.QUIZ and self.option_id == self.poll.correct_option_id:
+            details.append("correct vote")
+        details = ", ".join(details)
+        details = "" if not details else f"""<span class="details">{details}</span>"""
+
+        return f"""
+        <div class="answer">- {self.option.text} {details}</div>
+        """
+
+
+class Poll(HtmlMedia):
+    def __init__(self, *args, message: PyroMessage):
+        self.poll = message.poll
+
+    def no_media(self) -> str:
+        return ""
+
+    def to_html(self) -> str:
+        poll_details = None
+        if self.poll.is_anonymous:
+            poll_details = "<div class=\"details\">Anonymous poll</div>"
+        if self.poll.is_closed:
+            poll_details = "<div class=\"details\">Closed poll</div>"
+
+        options = ""
+        for idx, opt in enumerate(self.poll.options):
+            options += PollOption(self.poll, idx).to_html()
+
+        return f"""
+        <div class="media_poll">
+            <div class="question bold">{self.poll.question}</div>
+
+            {"" if poll_details is None else poll_details}
+
+            {options}
+            
+            <div class="total details"{self.poll.total_voter_count} votes</div>
+        </div>
+        """

texport/media.py

@@ -6,8 +6,10 @@
 from .html.animation import Animation
 from .html.audio import Audio
 from .html.base import HtmlMedia
+from .html.contact import Contact
 from .html.document import Document
 from .html.photo import Photo
+from .html.poll import Poll
 from .html.sticker import Sticker
 from .html.video import Video
 from .html.video_note import VideoNote
@@ -19,6 +21,12 @@ class Media:
     attr_name: str
     dir_name: str
     html_media_type: type[HtmlMedia]
+    has_size_limit: bool = True
+    downloadable: bool = True
+
+    def __post_init__(self):
+        if not self.downloadable:
+            self.has_size_limit = False
 
     def get_media(self, message: PyroMessage):
         return getattr(message, self.attr_name)
@@ -33,4 +41,6 @@ def get_media(self, message: PyroMessage):
     MessageMediaType.VIDEO: Media("video", "videos", Video),
     MessageMediaType.ANIMATION: Media("animation", "animations", Animation),
     MessageMediaType.VIDEO_NOTE: Media("video_note", "video_notes", VideoNote),
+    MessageMediaType.POLL: Media("poll", "_", Poll, False, False),
+    MessageMediaType.CONTACT: Media("contact", "_", Contact, False, False),
 }