tg_reader.py

#!/usr/bin/env python


from __future__ import annotations
from types import TracebackType
from typing import Iterator, List, cast
from pyrogram.client import Client
import pyrogram.types as pt
from pyrogram.raw.functions.channels.get_forum_topics import GetForumTopics
from pyrogram.raw.functions.messages.read_discussion import ReadDiscussion
from pyrogram.raw.types.input_channel import InputChannel
from pyrogram.raw.types.input_peer_channel import InputPeerChannel
from pyrogram.raw.types.forum_topic import ForumTopic
from pyrogram.raw.types.messages.forum_topics import ForumTopics
from dataclasses import dataclass, field, InitVar
from datetime import datetime


@dataclass
class TgChat:
    """
    Чат или подчат (тема форума)
    """
    id: int = field(init=False)                     # ID чата
    name: str = field(init=False)                   # Название чата или темы
    topic_id: int | None = field(init=False)        # ID темы (None для обычных чатов)

    chat: InitVar[pt.Chat]                          # Чат (для темы - родительский форум)
    topic: InitVar[ForumTopic | None] = None        # Тема форума (для подчатов)

    @staticmethod
    def get_chat_name(chat: pt.Chat) -> str:
        """
        Извлекате удобочитаемое название чата.
        """
        if chat.title:
            return chat.title
        full_name = " ".join(
            part for part in (chat.first_name, chat.last_name) if part
        )
        return full_name or chat.username or str(chat.id)

    def __post_init__(self, chat: pt.Chat, topic: ForumTopic | None):
        self.id = chat.id
        if topic is None:
            self.name = TgChat.get_chat_name(chat)
            self.topic_id = None
        else:
            self.name = topic.title or str(topic.id)
            self.topic_id = topic.id

    def __str__(self):
        display_id = self.topic_id if self.topic_id is not None else self.id
        return f"{self.name}: {display_id}"


@dataclass
class TgMessage:
    """
    Сообщение чата с метаданными
    """
    id: int
    date: datetime
    author: str
    content: str
    chat_id: int

    def __str__(self):
        return f"[{self.date}] {self.author}: {self.content}"


class TgReader:
    """
    Класс предоставляет упрощённый API для доступа к чатам Telegram.
    """
    def __init__(self, api_id: int, api_hash: str):
        self.__client = Client(
            "telegram_account",
            api_id = api_id,
            api_hash = api_hash
        )
    
    def __enter__(self) -> 'TgReader':
        self.__client.start() # type: ignore[misc]
        return self
    
    def __exit__(
            self,
            exc_type: type[BaseException] | None,
            exc_value: BaseException | None,
            traceback: TracebackType | None
    ):
        self.__client.stop() # type: ignore[misc]

    def get_chats(self) -> list[TgChat]:
        """
        Получить список доступных чатов.
        """
        dialogs = cast(
            Iterator[pt.Dialog],
            self.__client.get_dialogs(),
        )
        return [TgChat(dialog.chat) for dialog in dialogs]
    
    @staticmethod
    def get_message_author(message: pt.Message) -> str:
        """
        Получить удобочитаемое имя автора сообщения.
        """
        if message.from_user:
            user = message.from_user

            full_name = " ".join(
                part for part in (user.first_name, user.last_name) if part
            )

            return full_name or user.username or str(user.id)
        elif message.sender_chat:
            chat = message.sender_chat
            return chat.title or chat.username or str(chat.id)
        else:
            return "<Unknown>"
    
    @staticmethod
    def get_message_content(message: pt.Message) -> str:
        """
        Получить в читаемом виде содержимое сообщения.
        """
        if message.text:
            return message.text
        elif message.caption:
            return message.caption
        elif message.media:
            return f"<медиа: {message.media.value}>"
        elif message.service:
            return f"<служебное событие: {message.service.value}>"
        else:
            return "<служебное событие>"

    @staticmethod
    def is_message_in_subtopic(message: pt.Message, topic_id: int) -> bool:
        """
        Проверить, принадлежит ли сообщение подчату (теме форума).
        """
        return (
            message.id == topic_id
            or message.reply_to_top_message_id == topic_id
            or message.reply_to_message_id == topic_id
        )

    
    def get_messages(
            self,
            chat_id: int,
            limit: int = 0,
            offset_date: datetime = datetime(2026, 7, 9),
            topic_id: int | None = None,
    ) -> list[TgMessage]:
        """
        Получить `limit` последних сообщений из чата `chat_id`.

        Parameters:
            chat_id (``int``):
                ID чата
            
            limit (``int``):
                Максимальное количество сообщений, по умолчанию - не ограниченно.

            offset_date (:py:obj:`~datetime.datetime`):
                Получаем сообщения, новее этой отметки времени.

            topic_id (:py:obj:`int`, *optional*):
                ID подчата (темы форума). Если указан, возвращаются
                сообщения только из этой темы.
        """
        history = cast(
            Iterator[pt.Message],
            self.__client.get_chat_history(
                chat_id=chat_id,
                # При фильтрации по теме нельзя заранее запросить у API
                # `limit` сообщений: теме может принадлежать лишь их часть
                limit=0 if topic_id is not None else limit,
            )
        )
        messages = []
        for msg in history:
            if msg.date < offset_date:
                break
            if topic_id is not None and not TgReader.is_message_in_subtopic(msg, topic_id):
                continue
            messages.append(
                TgMessage(
                    msg.id,
                    msg.date,
                    TgReader.get_message_author(msg),
                    TgReader.get_message_content(msg),
                    msg.chat.id
                )
            )
            if topic_id is not None and limit and len(messages) >= limit:
                break
        return list(reversed(messages))

    def mark_as_read(
            self,
            chat_id: int,
            max_message_id: int,
            topic_id: int | None = None,
    ):
        """
        Пометить сообщения чата/темы как прочитанные.

        Parameters:
            chat_id (``int``):
                ID чата

            max_message_id (``int``):
                ID последнего прочитанного сообщения

            topic_id (:py:obj:`int`, *optional*):
                ID темы форума. Если указан, помечается как прочитанная
                только эта тема.
        """
        if topic_id is not None:
            peer = self.__client.resolve_peer(chat_id)  # type: ignore[misc]
            self.__client.invoke(  # type: ignore[misc]
                ReadDiscussion(
                    peer=peer,  # type: ignore[arg-type]
                    msg_id=topic_id,
                    read_max_id=max_message_id
                )
            )
        else:
            self.__client.read_chat_history(chat_id, max_id=max_message_id)  # type: ignore[misc]

    def get_chat_name(self, chat_id: int) -> str:
        return TgChat.get_chat_name(
            cast(
                pt.Chat,
                self.__client.get_chat(chat_id)
            )
        )

    def get_topic_name(self, chat_id: int | str, topic_id: int) -> str:
        """
        Получить название темы форума по её ID.
        """
        for subchat in self.get_subchats(chat_id):
            if subchat.topic_id == topic_id:
                return subchat.name
        return str(topic_id)

    def get_subchats(self, chat_id: int | str) -> list[TgChat]:
        """
        Получить список тем форума `chat_id`.

        Возвращает пустой список, если чат не поддерживает темы форума.
        """
        try:
            peer = self.__client.resolve_peer(chat_id)
        except Exception:
            return []

        if not isinstance(peer, InputPeerChannel):
            return []

        input_channel = InputChannel(
            channel_id=peer.channel_id,
            access_hash=peer.access_hash,
        )

        try:
            parent = cast(pt.Chat, self.__client.get_chat(chat_id))
        except Exception:
            return []

        subchats: list[TgChat] = []
        offset_topic = 0
        limit = 100
        try:
            while True:
                topics = cast(
                    List[ForumTopic],
                    cast(
                        ForumTopics,
                        self.__client.invoke(
                            GetForumTopics(
                                channel=input_channel,  # type: ignore[arg-type]
                                offset_date=0,
                                offset_id=0,
                                offset_topic=offset_topic,
                                limit=limit,
                            )
                        ),
                    ).topics,
                )
                if not topics:
                    break
                subchats.extend(
                    TgChat(parent, topic)
                    for topic in topics
                )
                if len(topics) < limit:
                    break
                offset_topic = topics[-1].id
        except Exception:
            return []
        return subchats


if __name__ == "__main__":
    import os

    # Читаем API_KEY
    from dotenv import load_dotenv
    load_dotenv()

    def create_TgReader() -> TgReader:
        return TgReader(
            api_id=int(os.environ["TG_API_ID"]),
            api_hash=os.environ["TG_API_HASH"]
        )

    # Парсер аргументов скрипта
    from argparse import ArgumentParser
    arg_parser = ArgumentParser(
        description="Справка по командам"
    )
    subparsers = arg_parser.add_subparsers(
        dest="command",
        required=True,
        help="Действие"
    )
    list_parser = subparsers.add_parser(
        "list",
        help="Показать список доступных чатов"
    )
    show_parser = subparsers.add_parser(
        "show",
        help="Показать сообщения из чата"
    )
    show_parser.add_argument(
        "chat_id",
        help="ID чата"
    )
    show_parser.add_argument(
        "count",
        type=int,
        nargs="?",
        default=0,
        help="Количество последних сообщений (по умолчанию, не ограничено)"
    )
    show_parser.add_argument(
        "--topic-id",
        type=int,
        default=None,
        help="ID подчата (темы форума) для фильтрации сообщений"
    )
    subchats_parser = subparsers.add_parser(
        "topics",
        help="Показать ID тем форума"
    )
    subchats_parser.add_argument(
        "chat_id",
        help="ID чата"
    )
    args = arg_parser.parse_args()

    if args.command == "list":
        # Запрашиваем список чатов
        tg = create_TgReader()
        with tg:
            chats = tg.get_chats()
        for chat in chats:
            print(chat)
    elif args.command == "show":
        # Запрашиваем сообщения из чата
        tg = create_TgReader()
        with tg:
            messages = tg.get_messages(
                args.chat_id,
                args.count,
                topic_id=args.topic_id,
            )
        for msg in messages:
            print(msg)
    elif args.command == "topics":
        # Запрашиваем подчаты
        tg = create_TgReader()
        with tg:
            subchats = tg.get_subchats(args.chat_id)
        if not subchats:
            print("Подчаты не найдены (чат не поддерживает темы форума)")
        for subchat in subchats:
            print(subchat)
    else:
        # Пользователь не указал действие
        # Выводим справку
        arg_parser.print_help()