Make edit_message_media work with new files

jsmnbom · jsmnbom · commit c1fbed2fe386 · 2018-08-12T11:51:40.000+02:00
diff --git a/telegram/__init__.py b/telegram/__init__.py
@@ -91,13 +91,13 @@
 from .webhookinfo import WebhookInfo
 from .games.gamehighscore import GameHighScore
 from .update import Update
+from .files.inputmedia import (InputMedia, InputMediaVideo, InputMediaPhoto, InputMediaAnimation,
+                               InputMediaAudio, InputMediaDocument)
 from .bot import Bot
 from .constants import (MAX_MESSAGE_LENGTH, MAX_CAPTION_LENGTH, SUPPORTED_WEBHOOK_PORTS,
                         MAX_FILESIZE_DOWNLOAD, MAX_FILESIZE_UPLOAD,
                         MAX_MESSAGES_PER_SECOND_PER_CHAT, MAX_MESSAGES_PER_SECOND,
                         MAX_MESSAGES_PER_MINUTE_PER_GROUP)
-from .files.inputmedia import (InputMedia, InputMediaVideo, InputMediaPhoto, InputMediaAnimation,
-                               InputMediaAudio, InputMediaDocument)
 from .version import __version__  # flake8: noqa
 
 __author__ = 'devs@python-telegram-bot.org'
diff --git a/telegram/bot.py b/telegram/bot.py
@@ -1805,7 +1805,7 @@ def edit_message_media(self,
                 Identifier of the sent message.
             inline_message_id (:obj:`str`, optional): Required if chat_id and message_id are not
                 specified. Identifier of the inline message.
-            media (:class:`telegram.InputMedia`): A JSON-serialized object for a new media content
+            media (:class:`telegram.InputMedia`): An object for a new media content
                 of the message.
             reply_markup (:class:`telegram.ReplyMarkup`, optional): Additional interface options. A
                 JSON-serialized object for an inline keyboard, custom reply keyboard, instructions
@@ -1823,7 +1823,7 @@ def edit_message_media(self,
 
         url = '{0}/editMessageMedia'.format(self.base_url)
 
-        data = {'media': media.to_dict()}
+        data = {'media': media}
 
         if chat_id:
             data['chat_id'] = chat_id
diff --git a/telegram/utils/request.py b/telegram/utils/request.py
@@ -41,7 +41,7 @@
                   "how to properly install.")
     raise
 
-from telegram import (InputFile, TelegramError)
+from telegram import (InputFile, TelegramError, InputMedia)
 from telegram.error import (Unauthorized, NetworkError, TimedOut, BadRequest, ChatMigrated,
                             RetryAfter, InvalidToken)
 
@@ -276,22 +276,31 @@ def post(self, url, data, timeout=None):
         files = False
 
         for key, val in data.copy().items():
-            if key == 'media':
-                media = []
-                for m in val:
-                    media.append(m.to_dict())
-                    if isinstance(m.media, InputFile):
-                        data[m.media.attach] = m.media.field_tuple
-                data[key] = json.dumps(media)
-                files = True
-
             if isinstance(val, InputFile):
+                # Convert the InputFile to urllib3 field format
                 data[key] = val.field_tuple
                 files = True
-
-            if isinstance(val, (float, int)):
+            elif isinstance(val, (float, int)):
+                # Urllib3 doesn't like floats it seems
                 data[key] = str(val)
+            elif key == 'media':
+                # One media or multiple
+                if isinstance(val, InputMedia):
+                    # Attach and set val to attached name
+                    data[key] = val.to_json()
+                    if isinstance(val.media, InputFile):
+                        data[val.media.attach] = val.media.field_tuple
+                else:
+                    # Attach and set val to attached name for all
+                    media = []
+                    for m in val:
+                        media.append(m.to_dict())
+                        if isinstance(m.media, InputFile):
+                            data[m.media.attach] = m.media.field_tuple
+                    data[key] = json.dumps(media)
+                files = True
 
+        # Use multipart upload if we're uploading files, otherwise use JSON
         if files:
             result = self._request_wrapper('POST', url, fields=data, **urlopen_kwargs)
         else:
diff --git a/tests/test_inputmedia.py b/tests/test_inputmedia.py
@@ -49,7 +49,6 @@ def input_media_video(class_thumb_file):
 def input_media_photo(class_thumb_file):
     return InputMediaPhoto(media=TestInputMediaPhoto.media,
                            caption=TestInputMediaPhoto.caption,
-                           thumb=class_thumb_file,
                            parse_mode=TestInputMediaPhoto.parse_mode)
 
 
@@ -144,7 +143,6 @@ def test_expected_values(self, input_media_photo):
         assert input_media_photo.media == self.media
         assert input_media_photo.caption == self.caption
         assert input_media_photo.parse_mode == self.parse_mode
-        assert isinstance(input_media_photo.thumb, InputFile)
 
     def test_to_dict(self, input_media_photo):
         input_media_photo_dict = input_media_photo.to_dict()
@@ -337,3 +335,13 @@ def test_edit_message_media(self, bot, chat_id, media_group):
         mid = messages[-1].message_id
         new_message = bot.edit_message_media(chat_id=cid, message_id=mid, media=media_group[0])
         assert isinstance(new_message, Message)
+
+    @flaky(3, 1)
+    @pytest.mark.timeout(10)
+    def test_edit_message_media_new_file(self, bot, chat_id, media_group, thumb_file):
+        messages = bot.send_media_group(chat_id, media_group)
+        cid = messages[-1].chat.id
+        mid = messages[-1].message_id
+        new_message = bot.edit_message_media(chat_id=cid, message_id=mid,
+                                             media=InputMediaPhoto(thumb_file))
+        assert isinstance(new_message, Message)
