Bring back the possibility to use strings as callback query data

delivrance · delivrance · commit 94de75f7144b · 2019-05-13T21:04:17.000+02:00
In case bytes (which is the type used by telegram) can't be successfully
decoded into strings, the raw bytes are presented instead of trying to
decode by ignoring/replacing errors.
diff --git a/pyrogram/client/methods/bots/request_callback_answer.py b/pyrogram/client/methods/bots/request_callback_answer.py
@@ -27,7 +27,7 @@ def request_callback_answer(
         self,
         chat_id: Union[int, str],
         message_id: int,
-        callback_data: bytes,
+        callback_data: Union[str, bytes],
         timeout: int = 10
     ):
         """Request a callback answer from bots.
@@ -42,7 +42,7 @@ def request_callback_answer(
             message_id (``int``):
                 The message id the inline keyboard is attached on.
 
-            callback_data (``bytes``):
+            callback_data (``str`` | ``bytes``):
                 Callback data associated with the inline button you want to get the answer from.
 
             timeout (``int``, *optional*):
@@ -56,11 +56,15 @@ def request_callback_answer(
             RPCError: In case of a Telegram RPC error.
             TimeoutError: In case the bot fails to answer within 10 seconds.
         """
+
+        # Telegram only wants bytes, but we are allowed to pass strings too.
+        data = bytes(callback_data, "utf-8") if isinstance(callback_data, str) else callback_data
+
         return self.send(
             functions.messages.GetBotCallbackAnswer(
                 peer=self.resolve_peer(chat_id),
                 msg_id=message_id,
-                data=callback_data
+                data=data
             ),
             retries=0,
             timeout=timeout
diff --git a/pyrogram/client/types/keyboards/callback_query.py b/pyrogram/client/types/keyboards/callback_query.py
@@ -18,6 +18,7 @@
 
 from base64 import b64encode
 from struct import pack
+from typing import Union
 
 import pyrogram
 from pyrogram.api import types
@@ -51,7 +52,7 @@ class CallbackQuery(PyrogramType, Update):
         inline_message_id (``str``):
             Identifier of the message sent via the bot in inline mode, that originated the query.
 
-        data (``bytes``, *optional*):
+        data (``str`` | ``bytes``, *optional*):
             Data associated with the callback button. Be aware that a bad client can send arbitrary data in this field.
 
         game_short_name (``str``, *optional*):
@@ -70,7 +71,7 @@ def __init__(
         chat_instance: str,
         message: "pyrogram.Message" = None,
         inline_message_id: str = None,
-        data: bytes = None,
+        data: Union[str, bytes] = None,
         game_short_name: str = None
     ):
         super().__init__(client)
@@ -80,7 +81,7 @@ def __init__(
         self.chat_instance = chat_instance
         self.message = message
         self.inline_message_id = inline_message_id
-        self.data = str(data, "utf-8")
+        self.data = data
         self.game_short_name = game_short_name
 
     @staticmethod
@@ -110,13 +111,20 @@ def _parse(client, callback_query, users) -> "CallbackQuery":
                 b"-_"
             ).decode().rstrip("=")
 
+        # Try to decode callback query data into string. If that fails, fallback to bytes instead of decoding by
+        # ignoring/replacing errors, this way, button clicks will still work.
+        try:
+            data = callback_query.data.decode()
+        except UnicodeDecodeError:
+            data = callback_query.data
+
         return CallbackQuery(
             id=str(callback_query.query_id),
             from_user=User._parse(client, users[callback_query.user_id]),
             message=message,
             inline_message_id=inline_message_id,
             chat_instance=str(callback_query.chat_instance),
-            data=callback_query.data,
+            data=data,
             game_short_name=callback_query.game_short_name,
             client=client
         )
diff --git a/pyrogram/client/types/keyboards/inline_keyboard_button.py b/pyrogram/client/types/keyboards/inline_keyboard_button.py
@@ -35,7 +35,7 @@ class InlineKeyboardButton(PyrogramType):
         text (``str``):
             Label text on the button.
 
-        callback_data (``bytes``, *optional*):
+        callback_data (``str`` | ``bytes``, *optional*):
             Data to be sent in a callback query to the bot when button is pressed, 1-64 bytes.
 
         url (``str``, *optional*):
@@ -75,7 +75,7 @@ def __init__(
 
         self.text = str(text)
         self.url = url
-        self.callback_data = bytes(callback_data, "utf-8") if isinstance(callback_data, str) else callback_data
+        self.callback_data = callback_data
         self.switch_inline_query = switch_inline_query
         self.switch_inline_query_current_chat = switch_inline_query_current_chat
         self.callback_game = callback_game
@@ -90,9 +90,16 @@ def read(o):
             )
 
         if isinstance(o, KeyboardButtonCallback):
+            # Try decode data to keep it as string, but if fails, fallback to bytes so we don't lose any information,
+            # instead of decoding by ignoring/replacing errors.
+            try:
+                data = o.data.decode()
+            except UnicodeDecodeError:
+                data = o.data
+
             return InlineKeyboardButton(
                 text=o.text,
-                callback_data=o.data
+                callback_data=data
             )
 
         if isinstance(o, KeyboardButtonSwitchInline):
@@ -115,7 +122,9 @@ def read(o):
 
     def write(self):
         if self.callback_data is not None:
-            return KeyboardButtonCallback(text=self.text, data=self.callback_data)
+            # Telegram only wants bytes, but we are allowed to pass strings too, for convenience.
+            data = bytes(self.callback_data, "utf-8") if isinstance(self.callback_data, str) else self.callback_data
+            return KeyboardButtonCallback(text=self.text, data=data)
 
         if self.url is not None:
             return KeyboardButtonUrl(text=self.text, url=self.url)
