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..8f46afc85 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): """ @@ -2721,54 +2705,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 +3187,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.