Add a custom Message class

telethon/tl/custom/__init__.py

@@ -1,3 +1,5 @@
 from .draft import Draft
 from .dialog import Dialog
 from .input_sized_file import InputSizedFile
+from .messagebutton import MessageButton
+from .message import Message

telethon/tl/custom/message.py

@@ -0,0 +1,197 @@
+from .. import types
+from ...extensions import markdown
+from ...utils import get_peer_id
+from .messagebutton import MessageButton
+
+
+class Message:
+    """
+    Custom class that encapsulates a message providing an abstraction to
+    easily access some commonly needed features (such as the markdown text
+    or the text for a given message entity).
+
+    Attributes:
+
+        original_message (:tl:`Message`):
+            The original :tl:`Message` object.
+
+        Any other attribute:
+            Attributes not described here are the same as those available
+            in the original :tl:`Message`.
+    """
+    def __init__(self, client, original, entities=None):
+        self.original_message = original
+        self.__getattribute__ = self.original_message.__getattribute__
+        self.__str__ = self.original_message.__str__
+        self.__repr__ = self.original_message.__repr__
+        self.stringify = self.original_message.stringify
+        self.to_dict = self.original_message.to_dict
+        self._client = client
+        self._text = None
+        self._reply_to = None
+        self._buttons = None
+        self._buttons_flat = []
+        self.from_user = entities[self.original_message.from_id]
+        self.chat = entities[get_peer_id(self.original_message.to_id)]
+
+    def __getattribute__(self, item):
+        return getattr(self.original_message, item)
+
+    @property
+    def client(self):
+        return self._client
+
+    @property
+    def text(self):
+        """
+        The message text, markdown-formatted.
+        """
+        if self._text is None:
+            if not self.original_message.entities:
+                return self.original_message.message
+            self._text = markdown.unparse(self.original_message.message,
+                                          self.original_message.entities or [])
+        return self._text
+
+    @property
+    def raw_text(self):
+        """
+        The raw message text, ignoring any formatting.
+        """
+        return self.original_message.message
+
+    @property
+    def buttons(self):
+        """
+        Returns a matrix (list of lists) containing all buttons of the message
+        as `telethon.tl.custom.messagebutton.MessageButton` instances.
+        """
+        if self._buttons is None and self.original_message.reply_markup:
+            if isinstance(self.original_message.reply_markup, (
+                    types.ReplyInlineMarkup, types.ReplyKeyboardMarkup)):
+                self._buttons = [[
+                    MessageButton(self._client, button, self.from_user,
+                                  self.chat, self.original_message.id)
+                    for button in row.buttons
+                ] for row in self.original_message.reply_markup.rows]
+                self._buttons_flat = [x for row in self._buttons for x in row]
+        return self._buttons
+
+    @property
+    def button_count(self):
+        """
+        Returns the total button count.
+        """
+        return len(self._buttons_flat) if self.buttons else 0
+
+    @property
+    def reply_to(self):
+        """
+        The :tl:`Message` that this message is replying to, or ``None``.
+
+        Note that this will make a network call to fetch the message and
+        will later be cached.
+        """
+        if self._reply_to is None:
+            if not self.original_message.reply_to_msg_id:
+                return None
+            self._reply_to = self._client.get_messages(
+                self.original_message.to_id,
+                ids=self.original_message.reply_to_msg_id
+            )
+
+    def reply(self, *args, **kwargs):
+        """
+        Replies to the message (as a reply). Shorthand for
+        `telethon.telegram_client.TelegramClient.send_message` with
+        both ``entity`` and ``reply_to`` already set.
+        """
+        kwargs['reply_to'] = self.original_message.id
+        return self._client.send_message(self.original_message.to_id,
+                                         *args, **kwargs)
+
+    def download_media(self, *args, **kwargs):
+        """
+        Downloads the media contained in the message, if any.
+        `telethon.telegram_client.TelegramClient.download_media` with
+        the ``message`` already set.
+        """
+        return self._client.download_media(self.original_message,
+                                           *args, **kwargs)
+
+    def get_entities_text(self):
+        """
+        Returns a list of tuples [(:tl:`MessageEntity`, `str`)], the string
+        being the inner text of the message entity (like bold, italics, etc).
+        """
+        texts = markdown.get_inner_text(self.original_message.message,
+                                        self.original_message.entities)
+        return list(zip(self.original_message.entities, texts))
+
+    def click(self, i=None, j=None, *, text=None, filter=None):
+        """
+        Clicks the inline keyboard button of the message, if any.
+
+        If the message has a non-inline keyboard, clicking it will
+        send the message, switch to inline, or open its URL.
+
+        Args:
+            i (`int`):
+                Clicks the i'th button (starting from the index 0).
+                Will ``raise IndexError`` if out of bounds. Example:
+
+                >>> message = Message(...)
+                >>> # Clicking the 3rd button
+                >>> # [button1] [button2]
+                >>> # [     button3     ]
+                >>> # [button4] [button5]
+                >>> message.click(2)  # index
+
+            j (`int`):
+                Clicks the button at position (i, j), these being the
+                indices for the (row, column) respectively. Example:
+
+                >>> # Clicking the 2nd button on the 1st row.
+                >>> # [button1] [button2]
+                >>> # [     button3     ]
+                >>> # [button4] [button5]
+                >>> message.click(0, 1)  # (row, column)
+
+                This is equivalent to ``message.buttons[0][1].click()``.
+
+            text (`str` | `callable`):
+                Clicks the first button with the text "text". This may
+                also be a callable, like a ``re.compile(...).match``,
+                and the text will be passed to it.
+
+            filter (`callable`):
+                Clicks the first button for which the callable
+                returns ``True``. The callable should accept a single
+                `telethon.tl.custom.messagebutton.MessageButton` argument.
+        """
+        if sum(int(x is not None) for x in (i, text, filter)) >= 2:
+            raise ValueError('You can only set either of i, text or filter')
+
+        if text is not None:
+            if callable(text):
+                for button in self._buttons_flat:
+                    if text(button.text):
+                        return button.click()
+            else:
+                for button in self._buttons_flat:
+                    if button.text == text:
+                        return button.click()
+            return
+
+        if filter is not None:
+            for button in self._buttons_flat:
+                if filter(button):
+                    return button.click()
+            return
+
+        if i is None:
+            i = 0
+        if j is None:
+            return self._buttons_flat[i].click()
+        else:
+            return self._buttons[i][j].click()

telethon/tl/custom/messagebutton.py

@@ -0,0 +1,69 @@
+from .. import types, functions
+import webbrowser
+
+
+class MessageButton:
+    """
+    Custom class that encapsulates a message providing an abstraction to
+    easily access some commonly needed features (such as the markdown text
+    or the text for a given message entity).
+
+    Attributes:
+
+        button (:tl:`KeyboardButton`):
+            The original :tl:`KeyboardButton` object.
+    """
+    def __init__(self, client, original, from_user, chat, msg_id):
+        self.button = original
+        self._from = from_user
+        self._chat = chat
+        self._msg_id = msg_id
+        self._client = client
+
+    @property
+    def client(self):
+        return self._client
+
+    @property
+    def text(self):
+        """The text string of the button."""
+        return self.button.text
+
+    @property
+    def data(self):
+        """The ``bytes`` data for :tl:`KeyboardButtonCallback` objects."""
+        if isinstance(self.button, types.KeyboardButtonCallback):
+            return self.button.data
+
+    @property
+    def inline_query(self):
+        """The query ``str`` for :tl:`KeyboardButtonSwitchInline` objects."""
+        if isinstance(self.button, types.KeyboardButtonSwitchInline):
+            return self.button.query
+
+    @property
+    def url(self):
+        """The url ``str`` for :tl:`KeyboardButtonUrl` objects."""
+        if isinstance(self.button, types.KeyboardButtonUrl):
+            return self.button.url
+
+    def click(self):
+        """
+        Clicks the inline keyboard button of the message, if any.
+
+        If the message has a non-inline keyboard, clicking it will
+        send the message, switch to inline, or open its URL.
+        """
+        if isinstance(self.button, types.KeyboardButton):
+            return self._client.send_message(
+                self._chat, self.button.text, reply_to=self._msg_id)
+        elif isinstance(self.button, types.KeyboardButtonCallback):
+            return self._client(functions.messages.GetBotCallbackAnswerRequest(
+                peer=self._chat, msg_id=self._msg_id, data=self.button.data
+            ), retries=1)
+        elif isinstance(self.button, types.KeyboardButtonSwitchInline):
+            return self._client(functions.messages.StartBotRequest(
+                bot=self._from, peer=self._chat, start_param=self.button.query
+            ), retries=1)
+        elif isinstance(self.button, types.KeyboardButtonUrl):
+            return webbrowser.open(self.button.url)