diff --git a/telebot/__init__.py b/telebot/__init__.py index e0c4c54..7756ab3 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -95,6 +95,9 @@ class TeleBot: self.num_threads = num_threads self.__create_threads = create_threads + self.message_subscribers_messages = [] + self.message_subscribers_callbacks = [] + self.message_handlers = [] if self.__create_threads: self.worker_pool = ThreadPool(num_threads) @@ -119,6 +122,7 @@ class TeleBot: def process_new_messages(self, new_messages): self.__notify_update(new_messages) self._notify_command_handlers(new_messages) + self._notify_message_subscribers(new_messages) def __notify_update(self, new_messages): for listener in self.update_listener: @@ -234,7 +238,7 @@ class TeleBot: def send_document(self, chat_id, data, reply_to_message_id=None, reply_markup=None): """ Use this method to send general files. - :param chat_id: + :param c.append()hat_id: :param data: :param reply_to_message_id: :param reply_markup: @@ -293,8 +297,41 @@ class TeleBot: return apihelper.send_chat_action(self.token, chat_id, action) def reply_to(self, message, text, **kwargs): + """ + Convenience function for `send_message(message.chat.id, text, reply_to_message_id=message.message_id, **kwargs)` + """ return self.send_message(message.chat.id, text, reply_to_message_id=message.message_id, **kwargs) + def register_for_reply(self, message, callback): + """ + Registers a callback function to be notified when a reply to `message` arrives. + + Warning: `message` must be sent with reply_markup=types.ForceReply(), otherwise TeleBot will not be able to see + the difference between a reply to `message` and an ordinary message. + + :param message: The message for which we are awaiting a reply. + :param callback: The callback function to be called when a reply arrives. Must accept one `message` + parameter, which will contain the replied message. + """ + self.message_subscribers_messages.insert(0, message.message_id) + self.message_subscribers_callbacks.insert(0, callback) + if len(self.message_subscribers_messages) > 10000: + self.message_subscribers_messages.pop() + self.message_subscribers_callbacks.pop() + + def _notify_message_subscribers(self, new_messages): + for message in new_messages: + if not hasattr(message, 'reply_to_message'): + continue + + reply_msg_id = message.reply_to_message.message_id + if reply_msg_id in self.message_subscribers_messages: + index = self.message_subscribers_messages.index(reply_msg_id) + self.message_subscribers_callbacks[index](message) + + del self.message_subscribers_messages[index] + del self.message_subscribers_callbacks[index] + def message_handler(self, commands=None, regexp=None, func=None, content_types=['text']): """ Message handler decorator. diff --git a/tests/test_telebot.py b/tests/test_telebot.py index 4f2f079..523d4d0 100644 --- a/tests/test_telebot.py +++ b/tests/test_telebot.py @@ -123,6 +123,18 @@ def test_reply_to(): ret_msg = tb.reply_to(msg, text + ' REPLY') assert ret_msg.reply_to_message.message_id == msg.message_id +def test_register_for_reply(): + text = 'CI reply_to Test Message' + tb = telebot.TeleBot(TOKEN) + msg = tb.send_message(CHAT_ID, text, reply_markup=types.ForceReply()) + reply_msg = tb.reply_to(msg, text + ' REPLY') + + def process_reply(message): + assert msg.message_id == message.reply_to_message.message_id + + tb.register_for_reply(msg, process_reply) + + tb.process_new_messages([reply_msg]) def test_send_location(): tb = telebot.TeleBot(TOKEN)