From 2e4280a94747f2572d3e78b2edef7929c79c1e71 Mon Sep 17 00:00:00 2001 From: _run Date: Fri, 1 Oct 2021 15:56:54 +0500 Subject: [PATCH 1/3] Update of state handlers No need to create state handlers --- examples/custom_states.py | 42 ++++++++++------ telebot/__init__.py | 97 +++++++------------------------------ telebot/custom_filters.py | 27 +++++++++++ telebot/handler_backends.py | 4 ++ 4 files changed, 77 insertions(+), 93 deletions(-) diff --git a/examples/custom_states.py b/examples/custom_states.py index 309fa6cd0..1d46f6c19 100644 --- a/examples/custom_states.py +++ b/examples/custom_states.py @@ -1,34 +1,48 @@ import telebot -from telebot.handler_backends import State +from telebot import custom_filters bot = telebot.TeleBot("") + + @bot.message_handler(commands=['start']) def start_ex(message): bot.set_state(message.chat.id, 1) bot.send_message(message.chat.id, 'Hi, write me a name') + -@bot.state_handler(state=1) -def name_get(message, state:State): +@bot.message_handler(state="*", commands='cancel') +def any_state(message): + bot.send_message(message.chat.id, "Your state was cancelled.") + bot.delete_state(message.chat.id) + +@bot.message_handler(state=1) +def name_get(message): bot.send_message(message.chat.id, f'Now write me a surname') - state.set(message.chat.id, 2) - with state.retrieve_data(message.chat.id) as data: + bot.set_state(message.chat.id, 2) + with bot.retrieve_data(message.chat.id) as data: data['name'] = message.text -@bot.state_handler(state=2) -def ask_age(message, state:State): +@bot.message_handler(state=2) +def ask_age(message): bot.send_message(message.chat.id, "What is your age?") - state.set(message.chat.id, 3) - with state.retrieve_data(message.chat.id) as data: + bot.set_state(message.chat.id, 3) + with bot.retrieve_data(message.chat.id) as data: data['surname'] = message.text -@bot.state_handler(state=3) -def ready_for_answer(message, state: State): - with state.retrieve_data(message.chat.id) as data: +@bot.message_handler(state=3, is_digit=True) +def ready_for_answer(message): + with bot.retrieve_data(message.chat.id) as data: bot.send_message(message.chat.id, "Ready, take a look:\nName: {name}\nSurname: {surname}\nAge: {age}".format(name=data['name'], surname=data['surname'], age=message.text), parse_mode="html") - state.finish(message.chat.id) + bot.delete_state(message.chat.id) + +@bot.message_handler(state=3, is_digit=False) +def age_incorrect(message): + bot.send_message(message.chat.id, 'Looks like you are submitting a string in the field age. Please enter a number') -bot.infinity_polling() +bot.add_custom_filter(custom_filters.StateFilter(bot)) +bot.add_custom_filter(custom_filters.IsDigitFilter()) +bot.infinity_polling(skip_pending=True) \ No newline at end of file diff --git a/telebot/__init__.py b/telebot/__init__.py index bfefd9273..d3cc2e500 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -499,7 +499,6 @@ def process_new_updates(self, updates): def process_new_messages(self, new_messages): self._notify_next_handlers(new_messages) self._notify_reply_handlers(new_messages) - self._notify_state_handlers(new_messages) self.__notify_update(new_messages) self._notify_command_handlers(self.message_handlers, new_messages) @@ -2386,6 +2385,9 @@ def delete_state(self, chat_id): """ self.current_states.delete_state(chat_id) + def retrieve_data(self, chat_id): + return self.current_states.retrieve_data(chat_id) + def get_state(self, chat_id): """ Get current state of a user. @@ -2394,6 +2396,14 @@ def get_state(self, chat_id): """ return self.current_states.current_state(chat_id) + def add_data(self, chat_id, **kwargs): + """ + Add data to states. + :param chat_id: + """ + for key, value in kwargs.items(): + self.current_states._add_data(chat_id, key, value) + def register_next_step_handler_by_chat_id( self, chat_id: Union[int, str], callback: Callable, *args, **kwargs) -> None: """ @@ -2459,32 +2469,6 @@ def _notify_next_handlers(self, new_messages): new_messages.pop(i) # removing message that was detected with next_step_handler - def _notify_state_handlers(self, new_messages): - """ - Description: TBD - :param new_messages: - :return: - """ - if not self.current_states: return - - for i, message in enumerate(new_messages): - need_pop = False - user_state = self.current_states.current_state(message.from_user.id) - if user_state: - for handler in self.state_handlers: - if handler['filters']['state'] == user_state: - for message_filter, filter_value in handler['filters'].items(): - if filter_value is None: - continue - if not self._test_filter(message_filter, filter_value, message): - return False - need_pop = True - state = self.current_states - self._exec_task(handler["function"], message, state) - if need_pop: - new_messages.pop(i) # removing message that was detected by states - - @staticmethod def _build_handler_dict(handler, **filters): """ @@ -2548,7 +2532,7 @@ def add_middleware_handler(self, handler, update_types=None): else: self.default_middleware_handlers.append(handler) - def message_handler(self, commands=None, regexp=None, func=None, content_types=None, chat_types=None, **kwargs): + def message_handler(self, commands=None, regexp=None, func=None, content_types=None, chat_types=None, state=None, **kwargs): """ Message handler decorator. This decorator can be used to decorate functions that must handle certain types of messages. @@ -2591,6 +2575,9 @@ def default_command(message): if content_types is None: content_types = ["text"] + if type(state) is not list and state is not None: + state = [state] + if isinstance(commands, str): logger.warning("message_handler: 'commands' filter should be List of strings (commands), not string.") commands = [commands] @@ -2605,6 +2592,7 @@ def decorator(handler): content_types=content_types, commands=commands, regexp=regexp, + state=state, func=func, **kwargs) self.add_message_handler(handler_dict) @@ -2620,7 +2608,7 @@ def add_message_handler(self, handler_dict): """ self.message_handlers.append(handler_dict) - def register_message_handler(self, callback, content_types=None, commands=None, regexp=None, func=None, chat_types=None, **kwargs): + def register_message_handler(self, callback, content_types=None, commands=None, regexp=None, func=None, chat_types=None, state=None, **kwargs): """ Registers message handler. :param callback: function to be called @@ -2644,6 +2632,7 @@ def register_message_handler(self, callback, content_types=None, commands=None, content_types=content_types, commands=commands, regexp=regexp, + state=state, func=func, **kwargs) self.add_message_handler(handler_dict) @@ -2721,54 +2710,6 @@ def register_edited_message_handler(self, callback, content_types=None, commands self.add_edited_message_handler(handler_dict) - def state_handler(self, state, content_types=None, regexp=None, func=None, chat_types=None, **kwargs): - """ - State handler for getting input from a user. - :param state: state of a user - :param content_types: - :param regexp: - :param func: - :param chat_types: - """ - def decorator(handler): - handler_dict = self._build_handler_dict(handler, - state=state, - content_types=content_types, - regexp=regexp, - chat_types=chat_types, - func=func, - **kwargs) - self.add_state_handler(handler_dict) - return handler - - return decorator - - def add_state_handler(self, handler_dict): - """ - Adds the edit message handler - :param handler_dict: - :return: - """ - self.state_handlers.append(handler_dict) - - def register_state_handler(self, callback, state, content_types=None, regexp=None, func=None, chat_types=None, **kwargs): - """ - Register a state handler. - :param callback: function to be called - :param state: state to be checked - :param content_types: - :param func: - """ - handler_dict = self._build_handler_dict(callback=callback, - state=state, - content_types=content_types, - regexp=regexp, - chat_types=chat_types, - func=func, - **kwargs) - self.add_state_handler(handler_dict) - - def channel_post_handler(self, commands=None, regexp=None, func=None, content_types=None, **kwargs): """ Channel post handler decorator @@ -3251,8 +3192,6 @@ def _test_filter(self, message_filter, filter_value, message): return filter_value(message) elif self.custom_filters and message_filter in self.custom_filters: return self._check_filter(message_filter,filter_value,message) - elif message_filter == 'state': - return True else: return False diff --git a/telebot/custom_filters.py b/telebot/custom_filters.py index 2f3b9b18c..20b42c382 100644 --- a/telebot/custom_filters.py +++ b/telebot/custom_filters.py @@ -146,3 +146,30 @@ def __init__(self, bot): def check(self, message): return self._bot.get_chat_member(message.chat.id, message.from_user.id).status in ['creator', 'administrator'] +class StateFilter(AdvancedCustomFilter): + """ + Filter to check state. + + Example: + @bot.message_handler(state=1) + """ + def __init__(self, bot): + self.bot = bot + key = 'state' + + def check(self, message, text): + if self.bot.current_states.current_state(message.from_user.id) is False:return False + elif '*' in text:return True + return self.bot.current_states.current_state(message.from_user.id) in text + +class IsDigitFilter(SimpleCustomFilter): + """ + Filter to check whether the string is made up of only digits. + + Example: + @bot.message_handler(is_digit=True) + """ + key = 'is_digit' + + def check(self, message): + return message.text.isdigit() diff --git a/telebot/handler_backends.py b/telebot/handler_backends.py index 9a37fad6d..d695b8212 100644 --- a/telebot/handler_backends.py +++ b/telebot/handler_backends.py @@ -179,6 +179,10 @@ def set(self, chat_id, new_state): """ self.add_state(chat_id,new_state) + def _add_data(self, chat_id, key, value): + result = self._states[chat_id]['data'][key] = value + return result + def finish(self, chat_id): """ Finish(delete) state of a user. From ff35f2521161651631876f1b766111d5b893fe6f Mon Sep 17 00:00:00 2001 From: _run Date: Fri, 1 Oct 2021 16:08:01 +0500 Subject: [PATCH 2/3] Update __init__.py --- telebot/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/telebot/__init__.py b/telebot/__init__.py index d3cc2e500..de6d76a82 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -2532,7 +2532,7 @@ def add_middleware_handler(self, handler, update_types=None): else: self.default_middleware_handlers.append(handler) - def message_handler(self, commands=None, regexp=None, func=None, content_types=None, chat_types=None, state=None, **kwargs): + def message_handler(self, commands=None, regexp=None, func=None, content_types=None, chat_types=None, **kwargs): """ Message handler decorator. This decorator can be used to decorate functions that must handle certain types of messages. From f337abe06eb20cf1938eebac4245af8835bf0d79 Mon Sep 17 00:00:00 2001 From: _run Date: Fri, 1 Oct 2021 16:09:20 +0500 Subject: [PATCH 3/3] Update __init__.py --- telebot/__init__.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/telebot/__init__.py b/telebot/__init__.py index de6d76a82..8f46afc85 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -2575,9 +2575,6 @@ def default_command(message): if content_types is None: content_types = ["text"] - if type(state) is not list and state is not None: - state = [state] - if isinstance(commands, str): logger.warning("message_handler: 'commands' filter should be List of strings (commands), not string.") commands = [commands] @@ -2592,7 +2589,6 @@ def decorator(handler): content_types=content_types, commands=commands, regexp=regexp, - state=state, func=func, **kwargs) self.add_message_handler(handler_dict) @@ -2608,7 +2604,7 @@ def add_message_handler(self, handler_dict): """ self.message_handlers.append(handler_dict) - def register_message_handler(self, callback, content_types=None, commands=None, regexp=None, func=None, chat_types=None, state=None, **kwargs): + def register_message_handler(self, callback, content_types=None, commands=None, regexp=None, func=None, chat_types=None, **kwargs): """ Registers message handler. :param callback: function to be called @@ -2632,7 +2628,6 @@ def register_message_handler(self, callback, content_types=None, commands=None, content_types=content_types, commands=commands, regexp=regexp, - state=state, func=func, **kwargs) self.add_message_handler(handler_dict)