From 70fb266eeaa579d77dfc702cfffb1ba6decafbfc Mon Sep 17 00:00:00 2001
From: apepenkov <39992738+apepenkov@users.noreply.github.com>
Date: Sat, 11 Nov 2023 21:29:42 +0300
Subject: [PATCH] Add ButtonCallback.get_message (#4250)

---
 .../src/telethon/_impl/client/client/users.py | 29 +++++++++++++++++--
 .../telethon/_impl/client/events/queries.py   | 23 ++++++++++++++-
 .../telethon/_impl/client/types/message.py    | 13 +++++++++
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/client/src/telethon/_impl/client/client/users.py b/client/src/telethon/_impl/client/client/users.py
index ca45676a..22c99ae6 100644
--- a/client/src/telethon/_impl/client/client/users.py
+++ b/client/src/telethon/_impl/client/client/users.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, List, Optional, Sequence
+from typing import TYPE_CHECKING, List, Optional, Sequence, Union
 
 from ...mtproto import RpcError
 from ...session import PackedChat, PackedType
@@ -117,7 +117,7 @@ async def get_chats(self: Client, chats: Sequence[ChatLike]) -> List[Chat]:
     ]
 
 
-async def resolve_to_packed(client: Client, chat: ChatLike) -> PackedChat:
+async def resolve_to_packed(client: Client, chat: Union[ChatLike, abcs.InputPeer, abcs.Peer]) -> PackedChat:
     if isinstance(chat, PackedChat):
         return chat
 
@@ -172,6 +172,31 @@ async def resolve_to_packed(client: Client, chat: ChatLike) -> PackedChat:
         else:
             raise RuntimeError("unexpected case")
 
+    if isinstance(chat, abcs.Peer):
+        packed = client._chat_hashes.get(peer_id(chat))
+        if packed is not None:
+            return packed
+        if isinstance(chat, types.PeerUser):
+            return PackedChat(
+                ty=PackedType.USER,
+                id=chat.user_id,
+                access_hash=0,
+            )
+        elif isinstance(chat, types.PeerChat):
+            return PackedChat(
+                ty=PackedType.CHAT,
+                id=chat.chat_id,
+                access_hash=0,
+            )
+        elif isinstance(chat, types.PeerChannel):
+            return PackedChat(
+                ty=PackedType.BROADCAST,
+                id=chat.channel_id,
+                access_hash=0,
+            )
+        else:
+            raise RuntimeError("unexpected case")
+
     if isinstance(chat, str):
         if chat.startswith("+"):
             resolved = await resolve_phone(client, chat)
diff --git a/client/src/telethon/_impl/client/events/queries.py b/client/src/telethon/_impl/client/events/queries.py
index 34a4d423..074cf812 100644
--- a/client/src/telethon/_impl/client/events/queries.py
+++ b/client/src/telethon/_impl/client/events/queries.py
@@ -3,8 +3,10 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Dict, Optional, Self
 
 from ...tl import abcs, functions, types
-from ..types import Chat
+from ..types import Chat, Message
 from .event import Event
+from ..types.chat import peer_id
+from ..client.messages import CherryPickedList
 
 if TYPE_CHECKING:
     from ..client.client import Client
@@ -69,6 +71,25 @@ class ButtonCallback(Event):
             )
         )
 
+    async def get_message(self) -> Optional[Message]:
+        """
+        Get the :class:`~telethon.types.Message` containing the button that was clicked.
+
+        If the message is too old and is no longer accessible, :data:`None` is returned instead.
+        """
+
+        pid = peer_id(self._raw.peer)
+        chat = self._chat_map.get(pid)
+        if not chat:
+            chat = await self._client._resolve_to_packed(pid)
+
+        lst = CherryPickedList(self._client, chat, [])
+        lst._ids.append(types.InputMessageCallbackQuery(id=self._raw.msg_id, query_id=self._raw.query_id))
+
+        message = (await lst)[0]
+
+        return message or None
+
 
 class InlineQuery(Event):
     """
diff --git a/client/src/telethon/_impl/client/types/message.py b/client/src/telethon/_impl/client/types/message.py
index 409cdf45..8bfa94e4 100644
--- a/client/src/telethon/_impl/client/types/message.py
+++ b/client/src/telethon/_impl/client/types/message.py
@@ -45,6 +45,16 @@ class Message(metaclass=NoPublicConstructor):
 
     You can get a message from :class:`telethon.events.NewMessage`,
     or from methods such as :meth:`telethon.Client.get_messages`.
+
+    Empty messages can occur very rarely when fetching the message history.
+    In these cases, only the :attr:`id` and :attr`peer` properties are guaranteed to be present.
+    To determine whether a message is empty, its truthy value can be checked via :meth:`object.__bool__`:
+
+    .. code-block:: python
+
+        async for message in client.iter_messages(chat):
+            if not message:
+                print('Found empty message with ID', message.id)
     """
 
     def __init__(
@@ -487,6 +497,9 @@ class Message(metaclass=NoPublicConstructor):
         else:
             return False
 
+    def __bool__(self):
+        return not isinstance(self._raw, types.MessageEmpty)
+
 
 def build_msg_map(
     client: Client, messages: List[abcs.Message], chat_map: Dict[int, Chat]