From b91cce15ba11be042c74719aaf6f1916f6c7d6cf Mon Sep 17 00:00:00 2001
From: Nick80835 <nick80835@gmail.com>
Date: Mon, 15 Jun 2020 14:23:24 -0400
Subject: [PATCH] add maybe broken support for callback queries

---
 ubot/command_handler.py | 19 +++++++++++++++++++
 ubot/loader.py          | 14 ++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/ubot/command_handler.py b/ubot/command_handler.py
index 446e1a1..4959337 100644
--- a/ubot/command_handler.py
+++ b/ubot/command_handler.py
@@ -16,10 +16,12 @@ class CommandHandler():
         self.incoming_commands = {}
         self.inline_photo_commands = {}
         self.inline_article_commands = {}
+        self.callback_queries = {}
         self.logger = logger
         self.settings = settings
         client.add_event_handler(self.handle_incoming, events.NewMessage(incoming=True))
         client.add_event_handler(self.handle_inline, events.InlineQuery())
+        client.add_event_handler(self.handle_callback_query, events.CallbackQuery())
 
     async def handle_incoming(self, event):
         prefix = escape(self.settings.get_config("cmd_prefix") or '.')
@@ -136,6 +138,23 @@ class CommandHandler():
         except:
             pass
 
+    async def handle_callback_query(self, event):
+        data_str = event.data.decode("utf-8")
+        data_id = data_str.split("*")[0]
+        data_data = data_str.lstrip(data_id + "*")
+
+        for key, value in self.callback_queries.items():
+            if key == data_id:
+                event.args = data_data
+                event.extras = value["extras"]
+
+                try:
+                    await value["function"](event)
+                except Exception as exception:
+                    self.logger.warn(f"{value['function'].__name__} - {exception}")
+                    await event.reply(f"`An error occurred in {value['function'].__name__}: {exception}`")
+                    raise exception
+
     async def fallback_inline(self, event):
         defaults_dict = {**self.inline_photo_commands, **self.inline_article_commands}
 
diff --git a/ubot/loader.py b/ubot/loader.py
index cc37f43..4910b21 100644
--- a/ubot/loader.py
+++ b/ubot/loader.py
@@ -35,6 +35,7 @@ class Loader():
         self.command_handler.incoming_commands = {}
         self.command_handler.inline_photo_commands = {}
         self.command_handler.inline_article_commands = {}
+        self.command_handler.callback_queries = {}
         self.help_dict = {}
 
         errors = ""
@@ -99,6 +100,19 @@ class Loader():
 
         return decorator
 
+    def add_callback_query(self, data_id=None, **args):
+        data_id = args.get("data_id", data_id)
+
+        def decorator(func):
+            self.command_handler.callback_queries[data_id] = {
+                "function": func,
+                "extras": args.get('extras', None)
+            }
+
+            return func
+
+        return decorator
+
     async def get_text(self, event, with_reply=True, return_msg=False, default=None):
         if event.args:
             if return_msg: