diff --git a/examples/custom_states.py b/examples/custom_states.py new file mode 100644 index 0000000..9ad9b1c --- /dev/null +++ b/examples/custom_states.py @@ -0,0 +1,34 @@ +import telebot + +from telebot.handler_backends import State + +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.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: + data['name'] = message.text + + +@bot.state_handler(state=2) +def ask_age(message, state:State): + 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: + 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.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.polling() \ No newline at end of file diff --git a/examples/register_handler/main.py b/examples/register_handler.py similarity index 59% rename from examples/register_handler/main.py rename to examples/register_handler.py index 1d670e9..f733a3b 100644 --- a/examples/register_handler/main.py +++ b/examples/register_handler.py @@ -1,11 +1,13 @@ import telebot -from telebot import custom_filters -import config -bot = telebot.TeleBot(config.api_token) -import handlers +api_token = '1297441208:AAH5THRzLXvY5breGFzkrEOIj7zwCGzbQ-Y' -bot.register_message_handler(handlers.start_executor, commands=['start']) # Start command executor +bot = telebot.TeleBot(api_token) + +def start_executor(message): + bot.send_message(message.chat.id, 'Hello!') + +bot.register_message_handler(start_executor, commands=['start']) # Start command executor # See also # bot.register_callback_query_handler(*args, **kwargs) diff --git a/examples/register_handler/config.py b/examples/register_handler/config.py deleted file mode 100644 index 50f8a62..0000000 --- a/examples/register_handler/config.py +++ /dev/null @@ -1,5 +0,0 @@ -import telebot - -api_token = '' - -bot = telebot.TeleBot(api_token) \ No newline at end of file diff --git a/examples/register_handler/handlers.py b/examples/register_handler/handlers.py deleted file mode 100644 index 5d65b35..0000000 --- a/examples/register_handler/handlers.py +++ /dev/null @@ -1,9 +0,0 @@ -# All handlers can be written in this file -from config import bot - -def start_executor(message): - bot.send_message(message.chat.id, 'Hello!') - -# Write more handlers here if you wish. You don't need a decorator - -# Just create function and register in main file. \ No newline at end of file diff --git a/telebot/__init__.py b/telebot/__init__.py index 32f7b1d..fe0d256 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -27,7 +27,7 @@ logger.addHandler(console_output_handler) logger.setLevel(logging.ERROR) from telebot import apihelper, util, types -from telebot.handler_backends import MemoryHandlerBackend, FileHandlerBackend +from telebot.handler_backends import MemoryHandlerBackend, FileHandlerBackend, StateMachine, State REPLY_MARKUP_TYPES = Union[ @@ -186,6 +186,9 @@ class TeleBot: self.my_chat_member_handlers = [] self.chat_member_handlers = [] self.custom_filters = {} + self.state_handlers = [] + + self.current_states = StateMachine() if apihelper.ENABLE_MIDDLEWARE: self.typed_middleware_handlers = { @@ -495,6 +498,7 @@ class TeleBot: def process_new_messages(self, new_messages): self._notify_next_handlers(new_messages) + self._notify_state_handlers(new_messages) self._notify_reply_handlers(new_messages) self.__notify_update(new_messages) self._notify_command_handlers(self.message_handlers, new_messages) @@ -2362,6 +2366,31 @@ class TeleBot: chat_id = message.chat.id self.register_next_step_handler_by_chat_id(chat_id, callback, *args, **kwargs) + def set_state(self, chat_id, state): + """ + Sets a new state of a user. + :param chat_id: + :param state: new state. can be string or integer. + """ + self.current_states.add_state(chat_id, state) + + def delete_state(self, chat_id): + """ + Delete the current state of a user. + :param chat_id: + :return: + """ + self.current_states.delete_state(chat_id) + + def get_state(self, chat_id): + """ + Get current state of a user. + :param chat_id: + :return: state of a user + """ + return self.current_states.current_state(chat_id) + + def register_next_step_handler_by_chat_id( self, chat_id: Union[int, str], callback: Callable, *args, **kwargs) -> None: """ @@ -2426,6 +2455,26 @@ class TeleBot: if need_pop: 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: + """ + 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: + need_pop = True + state = 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): """ @@ -2661,6 +2710,44 @@ class TeleBot: **kwargs) self.add_edited_message_handler(handler_dict) + + def state_handler(self, state=None, content_types=None, func=None,**kwargs): + + def decorator(handler): + handler_dict = self._build_handler_dict(handler, + state=state, + content_types=content_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=None, content_types=None, func=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, + state=state, + content_types=content_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 diff --git a/telebot/handler_backends.py b/telebot/handler_backends.py index 9b54f7c..721aba6 100644 --- a/telebot/handler_backends.py +++ b/telebot/handler_backends.py @@ -141,3 +141,79 @@ class RedisHandlerBackend(HandlerBackend): handlers = pickle.loads(value) self.clear_handlers(handler_group_id) return handlers + + + +class StateMachine: + def __init__(self): + self._states = {} + + def add_state(self, chat_id, state): + """ + Add a state. + """ + if chat_id in self._states: + + self._states[int(chat_id)]['state'] = state + else: + self._states[chat_id] = {'state': state,'data': {}} + + def current_state(self, chat_id): + """Current state""" + if chat_id in self._states: return self._states[chat_id]['state'] + else: return False + + def delete_state(self, chat_id): + """Delete a state""" + return self._states.pop(chat_id) + + +class State: + """ + Base class for state managing. + """ + def __init__(self, obj: StateMachine) -> None: + self.obj = obj + + def set(self, chat_id, new_state): + """ + Set a new state for a user. + :param chat_id: + :param new_state: new_state of a user + """ + self.obj._states[chat_id]['state'] = new_state + + def finish(self, chat_id): + """ + Finish(delete) state of a user. + :param chat_id: + """ + self.obj._states.pop(chat_id) + + def retrieve_data(self, chat_id): + """ + Save input text. + + Usage: + with state.retrieve_data(message.chat.id) as data: + data['name'] = message.text + + Also, at the end of your 'Form' you can get the name: + data['name'] + """ + return StateContext(self.obj, chat_id) + +class StateContext: + """ + Class for data. + """ + def __init__(self , obj: StateMachine, chat_id) -> None: + self.obj = obj + self.chat_id = chat_id + self.data = obj._states[chat_id]['data'] + + def __enter__(self): + return self.data + + def __exit__(self, exc_type, exc_val, exc_tb): + return \ No newline at end of file