From 98044d6faa8f8b34df46f596e4442b57d98c9434 Mon Sep 17 00:00:00 2001 From: _run Date: Wed, 13 Oct 2021 18:34:36 +0500 Subject: [PATCH 1/2] File support for states File support. Now states can be saved in pickle file --- examples/custom_states.py | 4 ++ telebot/__init__.py | 16 ++++- telebot/custom_filters.py | 6 +- telebot/handler_backends.py | 131 ++++++++++++++++++++++++++++++++++-- 4 files changed, 148 insertions(+), 9 deletions(-) diff --git a/examples/custom_states.py b/examples/custom_states.py index ac70bb9..c7acfe5 100644 --- a/examples/custom_states.py +++ b/examples/custom_states.py @@ -61,4 +61,8 @@ def age_incorrect(message): bot.add_custom_filter(custom_filters.StateFilter(bot)) bot.add_custom_filter(custom_filters.IsDigitFilter()) + +# set saving states into file. +bot.enable_saving_states() # you can delete this if you do not need to save states + bot.infinity_polling(skip_pending=True) \ No newline at end of file diff --git a/telebot/__init__.py b/telebot/__init__.py index 8f46afc..b8b0c6e 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, State +from telebot.handler_backends import MemoryHandlerBackend, FileHandlerBackend, StateMemory, StateFile REPLY_MARKUP_TYPES = Union[ @@ -188,7 +188,8 @@ class TeleBot: self.custom_filters = {} self.state_handlers = [] - self.current_states = State() + self.current_states = StateMemory() + if apihelper.ENABLE_MIDDLEWARE: self.typed_middleware_handlers = { @@ -237,6 +238,17 @@ class TeleBot: """ self.next_step_backend = FileHandlerBackend(self.next_step_backend.handlers, filename, delay) + def enable_saving_states(self, filename="./.state-save/states.pkl"): + """ + Enable saving states (by default saving disabled) + + :param filename: Filename of saving file + + """ + + self.current_states = StateFile(filename=filename) + self.current_states._create_dir() + def enable_save_reply_handlers(self, delay=120, filename="./.handler-saves/reply.save"): """ Enable saving reply handlers (by default saving disable) diff --git a/telebot/custom_filters.py b/telebot/custom_filters.py index bce2399..0b95523 100644 --- a/telebot/custom_filters.py +++ b/telebot/custom_filters.py @@ -158,9 +158,9 @@ class StateFilter(AdvancedCustomFilter): key = 'state' def check(self, message, text): - if self.bot.current_states.current_state(message.from_user.id) is False:return False - elif text == '*':return True - elif type(text) is list:return self.bot.current_states.current_state(message.from_user.id) in text + if self.bot.current_states.current_state(message.from_user.id) is False: return False + elif text == '*': return True + elif type(text) is list: return self.bot.current_states.current_state(message.from_user.id) in text return self.bot.current_states.current_state(message.from_user.id) == text class IsDigitFilter(SimpleCustomFilter): diff --git a/telebot/handler_backends.py b/telebot/handler_backends.py index d695b82..45b903b 100644 --- a/telebot/handler_backends.py +++ b/telebot/handler_backends.py @@ -143,7 +143,7 @@ class RedisHandlerBackend(HandlerBackend): return handlers -class State: +class StateMemory: def __init__(self): self._states = {} @@ -166,7 +166,7 @@ class State: def delete_state(self, chat_id): """Delete a state""" - return self._states.pop(chat_id) + self._states.pop(chat_id) def _get_data(self, chat_id): return self._states[chat_id]['data'] @@ -195,7 +195,7 @@ class State: Save input text. Usage: - with state.retrieve_data(message.chat.id) as data: + with bot.retrieve_data(message.chat.id) as data: data['name'] = message.text Also, at the end of your 'Form' you can get the name: @@ -203,11 +203,114 @@ class State: """ return StateContext(self, chat_id) + +class StateFile: + """ + Class to save states in a file. + """ + def __init__(self, filename): + self.file_path = filename + + def add_state(self, chat_id, state): + """ + Add a state. + :param chat_id: + :param state: new state + """ + states_data = self._read_data() + if chat_id in states_data: + states_data[chat_id]['state'] = state + return self._save_data(states_data) + else: + new_data = states_data[chat_id] = {'state': state,'data': {}} + return self._save_data(states_data) + + + def current_state(self, chat_id): + """Current state.""" + states_data = self._read_data() + if chat_id in states_data: return states_data[chat_id]['state'] + else: return False + + def delete_state(self, chat_id): + """Delete a state""" + states_data = self._read_data() + states_data.pop(chat_id) + self._save_data(states_data) + + def _read_data(self): + """ + Read the data from file. + """ + file = open(self.file_path, 'rb') + states_data = pickle.load(file) + file.close() + return states_data + + def _create_dir(self): + """ + Create directory .save-handlers. + """ + dirs = self.file_path.rsplit('/', maxsplit=1)[0] + os.makedirs(dirs, exist_ok=True) + if not os.path.isfile(self.file_path): + with open(self.file_path,'wb') as file: + pickle.dump({}, file) + + def _save_data(self, new_data): + """ + Save data after editing. + :param new_data: + """ + with open(self.file_path, 'wb+') as state_file: + pickle.dump(new_data, state_file, protocol=pickle.HIGHEST_PROTOCOL) + return True + + def _get_data(self, chat_id): + return self._read_data()[chat_id]['data'] + + 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.add_state(chat_id,new_state) + + def _add_data(self, chat_id, key, value): + states_data = self._read_data() + result = states_data[chat_id]['data'][key] = value + self._save_data(result) + + return result + + def finish(self, chat_id): + """ + Finish(delete) state of a user. + :param chat_id: + """ + self.delete_state(chat_id) + + def retrieve_data(self, chat_id): + """ + Save input text. + + Usage: + with bot.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 StateFileContext(self, chat_id) + + class StateContext: """ Class for data. """ - def __init__(self , obj: State, chat_id) -> None: + def __init__(self , obj: StateMemory, chat_id) -> None: self.obj = obj self.chat_id = chat_id self.data = obj._get_data(chat_id) @@ -217,3 +320,23 @@ class StateContext: def __exit__(self, exc_type, exc_val, exc_tb): return + +class StateFileContext: + """ + Class for data. + """ + def __init__(self , obj: StateFile, chat_id) -> None: + self.obj = obj + self.chat_id = chat_id + self.data = self.obj._get_data(self.chat_id) + + def __enter__(self): + return self.data + + def __exit__(self, exc_type, exc_val, exc_tb): + old_data = self.obj._read_data() + for i in self.data: + old_data[self.chat_id]['data'][i] = self.data.get(i) + self.obj._save_data(old_data) + + return From b6625baec63038cb2733c1ba5a951a3275713457 Mon Sep 17 00:00:00 2001 From: _run Date: Wed, 13 Oct 2021 19:02:17 +0500 Subject: [PATCH 2/2] Update __init__.py --- telebot/__init__.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/telebot/__init__.py b/telebot/__init__.py index b8b0c6e..64cf2c7 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -357,7 +357,7 @@ class TeleBot: """ return apihelper.delete_webhook(self.token, drop_pending_updates, timeout) - def get_webhook_info(self, timeout=None): + def get_webhook_info(self, timeout: Optional[int]=None): """ Use this method to get current webhook status. Requires no parameters. If the bot is using getUpdates, will return an object with the url field empty. @@ -2381,7 +2381,7 @@ 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): + def set_state(self, chat_id: int, state: Union[int, str]): """ Sets a new state of a user. :param chat_id: @@ -2389,7 +2389,7 @@ class TeleBot: """ self.current_states.add_state(chat_id, state) - def delete_state(self, chat_id): + def delete_state(self, chat_id: int): """ Delete the current state of a user. :param chat_id: @@ -2397,10 +2397,10 @@ class TeleBot: """ self.current_states.delete_state(chat_id) - def retrieve_data(self, chat_id): + def retrieve_data(self, chat_id: int): return self.current_states.retrieve_data(chat_id) - def get_state(self, chat_id): + def get_state(self, chat_id: int): """ Get current state of a user. :param chat_id: @@ -2408,7 +2408,7 @@ class TeleBot: """ return self.current_states.current_state(chat_id) - def add_data(self, chat_id, **kwargs): + def add_data(self, chat_id: int, **kwargs): """ Add data to states. :param chat_id: