From 296b866234ad0eb0ebccbc5d8e68958877302321 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sun, 24 Apr 2022 11:56:07 +0200
Subject: [PATCH] Improve performance by adding a message cache

---
 pyrogram/client.py                            | 21 +++++++++++++++
 .../bots_and_keyboards/callback_query.py      | 10 +++++--
 pyrogram/types/messages_and_media/message.py  | 26 +++++++++++++------
 3 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/pyrogram/client.py b/pyrogram/client.py
index 19ee0510..79168a93 100644
--- a/pyrogram/client.py
+++ b/pyrogram/client.py
@@ -270,6 +270,8 @@ class Client(Methods):
         # Username used for mentioned bot commands, e.g.: /start@usernamebot
         self.username = None
 
+        self.message_cache = Cache(10000)
+
         self.loop = asyncio.get_event_loop()
 
     def __enter__(self):
@@ -989,3 +991,22 @@ class Client(Methods):
 
     def guess_extension(self, mime_type: str) -> Optional[str]:
         return self.mimetypes.guess_extension(mime_type)
+
+
+class Cache:
+    def __init__(self, capacity: int):
+        self.capacity = capacity
+        self.store = {}
+
+    def __getitem__(self, key):
+        return self.store.get(key, None)
+
+    def __setitem__(self, key, value):
+        if key in self.store:
+            del self.store[key]
+
+        self.store[key] = value
+
+        if len(self.store) > self.capacity:
+            for _ in range(self.capacity // 2 + 1):
+                del self.store[next(iter(self.store))]
diff --git a/pyrogram/types/bots_and_keyboards/callback_query.py b/pyrogram/types/bots_and_keyboards/callback_query.py
index 8c96d6cc..dc4741b2 100644
--- a/pyrogram/types/bots_and_keyboards/callback_query.py
+++ b/pyrogram/types/bots_and_keyboards/callback_query.py
@@ -89,12 +89,18 @@ class CallbackQuery(Object, Update):
         self.matches = matches
 
     @staticmethod
-    async def _parse(client, callback_query, users) -> "CallbackQuery":
+    async def _parse(client: "pyrogram.Client", callback_query, users) -> "CallbackQuery":
         message = None
         inline_message_id = None
 
         if isinstance(callback_query, raw.types.UpdateBotCallbackQuery):
-            message = await client.get_messages(utils.get_peer_id(callback_query.peer), callback_query.msg_id)
+            chat_id = utils.get_peer_id(callback_query.peer)
+            message_id = callback_query.msg_id
+
+            message = client.message_cache[(chat_id, message_id)]
+
+            if not message:
+                message = await client.get_messages(chat_id, message_id)
         elif isinstance(callback_query, raw.types.UpdateInlineBotCallbackQuery):
             inline_message_id = b64encode(
                 pack(
diff --git a/pyrogram/types/messages_and_media/message.py b/pyrogram/types/messages_and_media/message.py
index 0c639e35..ec5eed51 100644
--- a/pyrogram/types/messages_and_media/message.py
+++ b/pyrogram/types/messages_and_media/message.py
@@ -445,7 +445,7 @@ class Message(Object, Update):
 
     @staticmethod
     async def _parse(
-        client,
+        client: "pyrogram.Client",
         message: raw.base.Message,
         users: dict,
         chats: dict,
@@ -592,6 +592,8 @@ class Message(Object, Update):
                     except MessageIdsEmpty:
                         pass
 
+            client.message_cache[(parsed_message.chat.id, parsed_message.id)] = parsed_message
+
             return parsed_message
 
         if isinstance(message, raw.types.Message):
@@ -802,18 +804,26 @@ class Message(Object, Update):
             )
 
             if message.reply_to:
+                parsed_message.reply_to_message_id = message.reply_to.reply_to_msg_id
+                parsed_message.reply_to_top_message_id = message.reply_to.reply_to_top_id
+
                 if replies:
                     try:
-                        parsed_message.reply_to_message = await client.get_messages(
-                            parsed_message.chat.id,
-                            reply_to_message_ids=message.id,
-                            replies=replies - 1
-                        )
+                        key = (parsed_message.chat.id, parsed_message.reply_to_message_id)
+                        reply_to_message = client.message_cache[key]
+
+                        if not reply_to_message:
+                            reply_to_message = await client.get_messages(
+                                parsed_message.chat.id,
+                                reply_to_message_ids=message.id,
+                                replies=replies - 1
+                            )
+
+                        parsed_message.reply_to_message = reply_to_message
                     except MessageIdsEmpty:
                         pass
 
-                parsed_message.reply_to_message_id = message.reply_to.reply_to_msg_id
-                parsed_message.reply_to_top_message_id = message.reply_to.reply_to_top_id
+            client.message_cache[(parsed_message.chat.id, parsed_message.id)] = parsed_message
 
             return parsed_message