Merge pull request python-telegram-bot#423 from python-telegram-bot/edit-class-method

jsmnbom · web-flow · commit a5671a8fb188 · 2016-10-15T23:03:43.000+02:00
Add edit_* class methods
diff --git a/telegram/callbackquery.py b/telegram/callbackquery.py
@@ -69,3 +69,52 @@ def to_dict(self):
     def answer(self, *args, **kwargs):
         """Shortcut for ``bot.answerCallbackQuery(update.callback_query.id, *args, **kwargs)``"""
         return self.bot.answerCallbackQuery(self.id, *args, **kwargs)
+
+    def edit_message_text(self, *args, **kwargs):
+        """
+        Shortcut for either ``bot.editMessageText(chat_id=update.callback_query.message.chat_id, \
+message_id=update.callback_query.message.message_id, \
+*args, **kwargs)``
+        or ``bot.editMessageText(inline_message_id=update.callback_query.inline_message_id, \
+*args, **kwargs)``
+        """
+        if self.inline_message_id:
+            return self.bot.edit_message_text(
+                inline_message_id=self.inline_message_id, *args, **kwargs)
+        else:
+            return self.bot.edit_message_text(
+                chat_id=self.message.chat_id, message_id=self.message.message_id, *args, **kwargs)
+
+    def edit_message_caption(self, *args, **kwargs):
+        """
+        Shortcut for either
+        ``bot.editMessageCaption(chat_id=update.callback_query.message.chat_id, \
+message_id=update.callback_query.message.message_id, \
+*args, **kwargs)``
+        or
+        ``bot.editMessageCaption(inline_message_id=update.callback_query.inline_message_id, \
+*args, **kwargs)``
+        """
+        if self.inline_message_id:
+            return self.bot.edit_message_caption(
+                inline_message_id=self.inline_message_id, *args, **kwargs)
+        else:
+            return self.bot.edit_message_caption(
+                chat_id=self.message.chat_id, message_id=self.message.message_id, *args, **kwargs)
+
+    def edit_message_reply_markup(self, *args, **kwargs):
+        """
+        Shortcut for either
+        ``bot.editMessageReplyMarkup(chat_id=update.callback_query.message.chat_id, \
+message_id=update.callback_query.message.message_id, \
+*args, **kwargs)``
+        or
+        ``bot.editMessageReplyMarkup(inline_message_id=update.callback_query.inline_message_id, \
+*args, **kwargs)``
+        """
+        if self.inline_message_id:
+            return self.bot.edit_message_reply_markup(
+                inline_message_id=self.inline_message_id, *args, **kwargs)
+        else:
+            return self.bot.edit_message_reply_markup(
+                chat_id=self.message.chat_id, message_id=self.message.message_id, *args, **kwargs)
diff --git a/telegram/message.py b/telegram/message.py
@@ -412,6 +412,48 @@ def forward(self, chat_id, disable_notification=False):
             disable_notification=disable_notification,
             message_id=self.message_id)
 
+    def edit_text(self, *args, **kwargs):
+        """
+        Shortcut for ``bot.editMessageText(chat_id=message.chat_id,
+                                       message_id=message.message_id,
+                                       *args, **kwargs)``
+
+        Note:
+            You can only edit messages that the bot sent itself,
+            therefore this method can only be used on the
+            return value of the bot.send_* family of methods.
+        """
+        return self.bot.edit_message_text(
+            chat_id=self.chat_id, message_id=self.message_id, *args, **kwargs)
+
+    def edit_caption(self, *args, **kwargs):
+        """
+        Shortcut for ``bot.editMessageCaption(chat_id=message.chat_id,
+                                              message_id=message.message_id,
+                                              *args, **kwargs)``
+
+        Note:
+            You can only edit messages that the bot sent itself,
+            therefore this method can only be used on the
+            return value of the bot.send_* family of methods.
+        """
+        return self.bot.edit_message_caption(
+            chat_id=self.chat_id, message_id=self.message_id, *args, **kwargs)
+
+    def edit_reply_markup(self, *args, **kwargs):
+        """
+        Shortcut for ``bot.editReplyMarkup(chat_id=message.chat_id,
+                                           message_id=message.message_id,
+                                           *args, **kwargs)``
+
+        Note:
+            You can only edit messages that the bot sent itself,
+            therefore this method can only be used on the
+            return value of the bot.send_* family of methods.
+        """
+        return self.bot.edit_message_caption(
+            chat_id=self.chat_id, message_id=self.message_id, *args, **kwargs)
+
     def parse_entity(self, entity):
         """
         Returns the text from a given :class:`telegram.MessageEntity`.
diff --git a/tests/test_message.py b/tests/test_message.py
@@ -76,6 +76,15 @@ def test_forward(self):
         self.assertTrue(self.is_json(message.to_json()))
         self.assertEqual(message.text, 'Testing class method')
 
+    @flaky(3, 1)
+    def test_edit_text(self):
+        """Test for Message.edit_text"""
+        message = self._bot.sendMessage(self._chat_id, '.')
+        message = message.edit_text('Testing class method')
+
+        self.assertTrue(self.is_json(message.to_json()))
+        self.assertEqual(message.text, 'Testing class method')
+
 
 if __name__ == '__main__':
     unittest.main()
