Merge pull request #558 from uburuntu/InputMedia

v.4.0: InputMediaAnimation, InputMediaAudio, InputMediaDocument, editMessageMedia

telebot/__init__.py

@@ -998,6 +998,12 @@ class TeleBot:
             return result
         return types.Message.de_json(result)
 
+    def edit_message_media(self, media, chat_id=None, message_id=None, inline_message_id=None, reply_markup=None):
+        result = apihelper.edit_message_media(self.token, media, chat_id, message_id, inline_message_id, reply_markup)
+        if type(result) == bool:  # if edit inline message return is bool not Message.
+            return result
+        return types.Message.de_json(result)
+
     def edit_message_reply_markup(self, chat_id=None, message_id=None, inline_message_id=None, reply_markup=None):
         result = apihelper.edit_message_reply_markup(self.token, chat_id, message_id, inline_message_id, reply_markup)
         if type(result) == bool:
@@ -1675,6 +1681,10 @@ class AsyncTeleBot(TeleBot):
     def edit_message_text(self, *args, **kwargs):
         return TeleBot.edit_message_text(self, *args, **kwargs)
 
+    @util.async_dec()
+    def edit_message_media(self, *args, **kwargs):
+        return TeleBot.edit_message_media(self, *args, **kwargs)
+
     @util.async_dec()
     def edit_message_reply_markup(self, *args, **kwargs):
         return TeleBot.edit_message_reply_markup(self, *args, **kwargs)

telebot/apihelper.py

@@ -264,7 +264,7 @@ def send_photo(token, chat_id, photo, caption=None, reply_to_message_id=None, re
 
 def send_media_group(token, chat_id, media, disable_notification=None, reply_to_message_id=None):
     method_url = r'sendMediaGroup'
-    media_json, files = _convert_input_media(media)
+    media_json, files = _convert_input_media_array(media)
     payload = {'chat_id': chat_id, 'media': media_json}
     if disable_notification:
         payload['disable_notification'] = disable_notification
@@ -638,6 +638,21 @@ def edit_message_caption(token, caption, chat_id=None, message_id=None, inline_m
     return _make_request(token, method_url, params=payload)
 
 
+def edit_message_media(token, media, chat_id=None, message_id=None, inline_message_id=None, reply_markup=None):
+    method_url = r'editMessageMedia'
+    media_json, file = _convert_input_media(media)
+    payload = {'media': media_json}
+    if chat_id:
+        payload['chat_id'] = chat_id
+    if message_id:
+        payload['message_id'] = message_id
+    if inline_message_id:
+        payload['inline_message_id'] = inline_message_id
+    if reply_markup:
+        payload['reply_markup'] = _convert_markup(reply_markup)
+    return _make_request(token, method_url, params=payload, files=file, method='post' if file else 'get')
+
+
 def edit_message_reply_markup(token, chat_id=None, message_id=None, inline_message_id=None, reply_markup=None):
     method_url = r'editMessageReplyMarkup'
     payload = {}
@@ -937,11 +952,17 @@ def _convert_markup(markup):
     return markup
 
 
-def _convert_input_media(array):
+def _convert_input_media(media):
+    if isinstance(media, types.InputMedia):
+        return media._convert_input_media()
+    return None, None
+
+
+def _convert_input_media_array(array):
     media = []
     files = {}
     for input_media in array:
-        if isinstance(input_media, types.JsonSerializable):
+        if isinstance(input_media, types.InputMedia):
             media_dict = input_media.to_dic()
             if media_dict['media'].startswith('attach://'):
                 key = media_dict['media'].replace('attach://', '')

telebot/types.py

@@ -2073,48 +2073,61 @@ class MaskPosition(JsonDeserializable, JsonSerializable):
 
 # InputMedia
 
-class InputMediaPhoto(JsonSerializable):
+class InputMedia(JsonSerializable):
-    def __init__(self, media, caption=None, parse_mode=None):
+    def __init__(self, type, media, caption=None, parse_mode=None):
-        self.type = "photo"
+        self.type = type
         self.media = media
         self.caption = caption
         self.parse_mode = parse_mode
 
+        if util.is_string(self.media):
+            self._media_name = ''
+            self._media_dic = self.media
+        else:
+            self._media_name = util.generate_random_token()
+            self._media_dic = 'attach://{}'.format(self._media_name)
+
     def to_json(self):
         return json.dumps(self.to_dic())
 
     def to_dic(self):
-        ret = {'type': self.type, 'media': 'attach://' + util.generate_random_token()
+        ret = {'type': self.type, 'media': self._media_dic}
-               if not util.is_string(self.media) else self.media}
         if self.caption:
             ret['caption'] = self.caption
         if self.parse_mode:
             ret['parse_mode'] = self.parse_mode
         return ret
 
+    def _convert_input_media(self):
+        if util.is_string(self.media):
+            return self.to_json(), None
 
-class InputMediaVideo(JsonSerializable):
+        return self.to_json(), {self._media_name: self.media}
-    def __init__(self, media, caption=None, parse_mode=None, width=None, height=None, duration=None,
+
+
+class InputMediaPhoto(InputMedia):
+    def __init__(self, media, caption=None, parse_mode=None):
+        super(InputMediaPhoto, self).__init__(type="photo", media=media, caption=caption, parse_mode=parse_mode)
+
+    def to_dic(self):
+        ret = super(InputMediaPhoto, self).to_dic()
+        return ret
+
+
+class InputMediaVideo(InputMedia):
+    def __init__(self, media, thumb=None, caption=None, parse_mode=None, width=None, height=None, duration=None,
                  supports_streaming=None):
-        self.type = "video"
+        super(InputMediaVideo, self).__init__(type="video", media=media, caption=caption, parse_mode=parse_mode)
-        self.media = media
+        self.thumb = thumb
-        self.caption = caption
-        self.parse_mode = parse_mode
         self.width = width
         self.height = height
         self.duration = duration
         self.supports_streaming = supports_streaming
 
-    def to_json(self):
-        return json.dumps(self.to_dic())
-
     def to_dic(self):
-        ret = {'type': self.type, 'media': 'attach://' + util.generate_random_token()
+        ret = super(InputMediaVideo, self).to_dic()
-               if not util.is_string(self.media) else self.media}
+        if self.thumb:
-        if self.caption:
+            ret['thumb'] = self.thumb
-            ret['caption'] = self.caption
-        if self.parse_mode:
-            ret['parse_mode'] = self.parse_mode
         if self.width:
             ret['width'] = self.width
         if self.height:
@@ -2124,3 +2137,57 @@ class InputMediaVideo(JsonSerializable):
         if self.supports_streaming:
             ret['supports_streaming'] = self.supports_streaming
         return ret
+
+
+class InputMediaAnimation(InputMedia):
+    def __init__(self, media, thumb=None, caption=None, parse_mode=None, width=None, height=None, duration=None):
+        super(InputMediaAnimation, self).__init__(type="animation", media=media, caption=caption, parse_mode=parse_mode)
+        self.thumb = thumb
+        self.width = width
+        self.height = height
+        self.duration = duration
+
+    def to_dic(self):
+        ret = super(InputMediaAnimation, self).to_dic()
+        if self.thumb:
+            ret['thumb'] = self.thumb
+        if self.width:
+            ret['width'] = self.width
+        if self.height:
+            ret['height'] = self.height
+        if self.duration:
+            ret['duration'] = self.duration
+        return ret
+
+
+class InputMediaAudio(InputMedia):
+    def __init__(self, media, thumb=None, caption=None, parse_mode=None, duration=None, performer=None, title=None):
+        super(InputMediaAudio, self).__init__(type="audio", media=media, caption=caption, parse_mode=parse_mode)
+        self.thumb = thumb
+        self.duration = duration
+        self.performer = performer
+        self.title = title
+
+    def to_dic(self):
+        ret = super(InputMediaAudio, self).to_dic()
+        if self.thumb:
+            ret['thumb'] = self.thumb
+        if self.duration:
+            ret['duration'] = self.duration
+        if self.performer:
+            ret['performer'] = self.performer
+        if self.title:
+            ret['title'] = self.title
+        return ret
+
+
+class InputMediaDocument(InputMedia):
+    def __init__(self, media, thumb=None, caption=None, parse_mode=None):
+        super(InputMediaDocument, self).__init__(type="document", media=media, caption=caption, parse_mode=parse_mode)
+        self.thumb = thumb
+
+    def to_dic(self):
+        ret = super(InputMediaDocument, self).to_dic()
+        if self.thumb:
+            ret['thumb'] = self.thumb
+        return ret

tests/test_telebot.py

@@ -361,6 +361,20 @@ class TestTeleBot:
         new_msg = tb.edit_message_caption(caption='Edit test', chat_id=CHAT_ID, message_id=msg.message_id)
         assert new_msg.caption == 'Edit test'
 
+    def test_edit_message_media(self):
+        file_data = open('../examples/detailed_example/kitten.jpg', 'rb')
+        file_data_2 = open('../examples/detailed_example/rooster.jpg', 'rb')
+        tb = telebot.TeleBot(TOKEN)
+        msg = tb.send_photo(CHAT_ID, file_data)
+        new_msg = tb.edit_message_media(chat_id=CHAT_ID, message_id=msg.message_id,
+                                        media=types.InputMediaPhoto(file_data_2, caption='Test editMessageMedia 0'))
+        assert type(new_msg) != bool
+
+        new_msg = tb.edit_message_media(chat_id=CHAT_ID, message_id=msg.message_id,
+                                        media=types.InputMediaPhoto(msg.photo[0].file_id, caption='Test editMessageMedia'))
+        assert type(new_msg) != bool
+        assert new_msg.caption == 'Test editMessageMedia'
+
     def test_get_chat(self):
         tb = telebot.TeleBot(TOKEN)
         ch = tb.get_chat(GROUP_ID)