From cf69a06ab858f75f8ba7fca6b1570ce512f0ba4a Mon Sep 17 00:00:00 2001 From: uburuntu Date: Fri, 10 Aug 2018 16:47:59 +0300 Subject: [PATCH] enh: make code better and enhance test case --- telebot/apihelper.py | 7 +++---- telebot/types.py | 29 +++++++++++++++++------------ tests/test_telebot.py | 7 +++++-- 3 files changed, 25 insertions(+), 18 deletions(-) diff --git a/telebot/apihelper.py b/telebot/apihelper.py index f74f7c7..ae72e37 100644 --- a/telebot/apihelper.py +++ b/telebot/apihelper.py @@ -1,5 +1,4 @@ # -*- coding: utf-8 -*- -from telebot.types import InputMedia try: import ujson as json @@ -651,7 +650,7 @@ def edit_message_media(token, media, chat_id=None, message_id=None, inline_messa 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) + 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): @@ -954,7 +953,7 @@ def _convert_markup(markup): def _convert_input_media(media): - if isinstance(media, InputMedia): + if isinstance(media, types.InputMedia): return media._convert_input_media() return None, None @@ -963,7 +962,7 @@ def _convert_input_media_array(array): media = [] files = {} for input_media in array: - if isinstance(input_media, InputMedia): + if isinstance(input_media, types.InputMedia): media_dict = input_media.to_dic() if media_dict['media'].startswith('attach://'): key = media_dict['media'].replace('attach://', '') diff --git a/telebot/types.py b/telebot/types.py index f3e0300..2e48b5d 100644 --- a/telebot/types.py +++ b/telebot/types.py @@ -2080,7 +2080,12 @@ class InputMedia(JsonSerializable): self.caption = caption self.parse_mode = parse_mode - self._media_dic = 'attach://' + util.generate_random_token() if not util.is_string(self.media) else self.media + 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()) @@ -2097,22 +2102,22 @@ class InputMedia(JsonSerializable): if util.is_string(self.media): return self.to_json(), None - return self.to_json(), {self._media_dic: self.media} + return self.to_json(), {self._media_name: self.media} class InputMediaPhoto(InputMedia): def __init__(self, media, caption=None, parse_mode=None): - super(InputMedia).__init__(type="photo", media=media, caption=caption, parse_mode=parse_mode) + super(InputMediaPhoto, self).__init__(type="photo", media=media, caption=caption, parse_mode=parse_mode) def to_dic(self): - ret = super(InputMedia).to_dic() + 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): - super(InputMedia).__init__(type="video", media=media, caption=caption, parse_mode=parse_mode) + super(InputMediaVideo, self).__init__(type="video", media=media, caption=caption, parse_mode=parse_mode) self.thumb = thumb self.width = width self.height = height @@ -2120,7 +2125,7 @@ class InputMediaVideo(InputMedia): self.supports_streaming = supports_streaming def to_dic(self): - ret = super(InputMedia).to_dic() + ret = super(InputMediaVideo, self).to_dic() if self.thumb: ret['thumb'] = self.thumb if self.width: @@ -2136,14 +2141,14 @@ class InputMediaVideo(InputMedia): class InputMediaAnimation(InputMedia): def __init__(self, media, thumb=None, caption=None, parse_mode=None, width=None, height=None, duration=None): - super(InputMedia).__init__(type="animation", media=media, caption=caption, parse_mode=parse_mode) + 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(InputMedia).to_dic() + ret = super(InputMediaAnimation, self).to_dic() if self.thumb: ret['thumb'] = self.thumb if self.width: @@ -2157,14 +2162,14 @@ class InputMediaAnimation(InputMedia): class InputMediaAudio(InputMedia): def __init__(self, media, thumb=None, caption=None, parse_mode=None, duration=None, performer=None, title=None): - super(InputMedia).__init__(type="audio", media=media, caption=caption, parse_mode=parse_mode) + 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(InputMedia).to_dic() + ret = super(InputMediaAudio, self).to_dic() if self.thumb: ret['thumb'] = self.thumb if self.duration: @@ -2178,11 +2183,11 @@ class InputMediaAudio(InputMedia): class InputMediaDocument(InputMedia): def __init__(self, media, thumb=None, caption=None, parse_mode=None): - super(InputMedia).__init__(type="document", media=media, caption=caption, parse_mode=parse_mode) + super(InputMediaDocument, self).__init__(type="document", media=media, caption=caption, parse_mode=parse_mode) self.thumb = thumb def to_dic(self): - ret = super(InputMedia).to_dic() + ret = super(InputMediaDocument, self).to_dic() if self.thumb: ret['thumb'] = self.thumb return ret diff --git a/tests/test_telebot.py b/tests/test_telebot.py index 3279b30..b81485e 100644 --- a/tests/test_telebot.py +++ b/tests/test_telebot.py @@ -366,11 +366,14 @@ class TestTeleBot: 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=file_data_2) + 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=msg.photo[0].file_id) + 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)