From 8045ad56ea84dcc6267e0ad800bcb4283c1e3a0d Mon Sep 17 00:00:00 2001 From: _run Date: Mon, 24 Jan 2022 21:24:56 +0400 Subject: [PATCH] States Update --- .../asynchronous_telebot/custom_states.py | 32 ++++++---- examples/custom_states.py | 58 +++++++++++++------ telebot/async_telebot.py | 17 +++--- telebot/asyncio_handler_backends.py | 16 +++++ telebot/handler_backends.py | 16 +++++ 5 files changed, 101 insertions(+), 38 deletions(-) diff --git a/examples/asynchronous_telebot/custom_states.py b/examples/asynchronous_telebot/custom_states.py index 56546bc..7658a81 100644 --- a/examples/asynchronous_telebot/custom_states.py +++ b/examples/asynchronous_telebot/custom_states.py @@ -1,14 +1,22 @@ import telebot from telebot import asyncio_filters from telebot.async_telebot import AsyncTeleBot + +# list of storages, you can use any storage +from telebot.asyncio_storage import StateRedisStorage,StateMemoryStorage,StatePickleStorage + +# new feature for states. +from telebot.asyncio_handler_backends import State, StatesGroup + bot = AsyncTeleBot('TOKEN') +# Just create different statesgroup +class MyStates(StatesGroup): + name = State() # statesgroup should contain states + surname = State() + age = State() -class MyStates: - name = 1 - surname = 2 - age = 3 @@ -17,7 +25,7 @@ async def start_ex(message): """ Start command. Here we are starting state """ - await bot.set_state(message.from_user.id, MyStates.name) + await bot.set_state(message.from_user.id, MyStates.name, message.chat.id) await bot.send_message(message.chat.id, 'Hi, write me a name') @@ -28,7 +36,7 @@ async def any_state(message): Cancel state """ await bot.send_message(message.chat.id, "Your state was cancelled.") - await bot.delete_state(message.from_user.id) + await bot.delete_state(message.from_user.id, message.chat.id) @bot.message_handler(state=MyStates.name) async def name_get(message): @@ -36,8 +44,8 @@ async def name_get(message): State 1. Will process when user's state is 1. """ await bot.send_message(message.chat.id, f'Now write me a surname') - await bot.set_state(message.from_user.id, MyStates.surname) - async with bot.retrieve_data(message.from_user.id) as data: + await bot.set_state(message.from_user.id, MyStates.surname, message.chat.id) + async with bot.retrieve_data(message.from_user.id, message.chat.id) as data: data['name'] = message.text @@ -47,16 +55,16 @@ async def ask_age(message): State 2. Will process when user's state is 2. """ await bot.send_message(message.chat.id, "What is your age?") - await bot.set_state(message.from_user.id, MyStates.age) - async with bot.retrieve_data(message.from_user.id) as data: + await bot.set_state(message.from_user.id, MyStates.age, message.chat.id) + async with bot.retrieve_data(message.from_user.id, message.chat.id) as data: data['surname'] = message.text # result @bot.message_handler(state=MyStates.age, is_digit=True) async def ready_for_answer(message): - async with bot.retrieve_data(message.from_user.id) as data: + async with bot.retrieve_data(message.from_user.id, message.chat.id) as data: await 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") - await bot.delete_state(message.from_user.id) + await bot.delete_state(message.from_user.id, message.chat.id) #incorrect number @bot.message_handler(state=MyStates.age, is_digit=False) diff --git a/examples/custom_states.py b/examples/custom_states.py index 5acc8f2..7901896 100644 --- a/examples/custom_states.py +++ b/examples/custom_states.py @@ -1,14 +1,39 @@ -import telebot +import telebot # telebot from telebot import custom_filters +from telebot.handler_backends import State, StatesGroup #States -bot = telebot.TeleBot("") +# States storage +from telebot.storage import StateRedisStorage, StatePickleStorage, StateMemoryStorage -class MyStates: - name = 1 - surname = 2 - age = 3 +# Beginning from version 4.4.0+, we support storages. +# StateRedisStorage -> Redis-based storage. +# StatePickleStorage -> Pickle-based storage. +# For redis, you will need to install redis. +# Pass host, db, password, or anything else, +# if you need to change config for redis. +# Pickle requires path. Default path is in folder .state-saves. +# If you were using older version of pytba for pickle, +# you need to migrate from old pickle to new by using +# StatePickleStorage().convert_old_to_new() + + + +# Now, you can pass storage to bot. +state_storage = StateMemoryStorage() # you can init here another storage + +bot = telebot.TeleBot("TOKEN", +state_storage=state_storage) + + +# States group. +class MyStates(StatesGroup): + # Just name variables differently + name = State() # creating instances of State class is enough from now + surname = State() + age = State() + @@ -17,18 +42,18 @@ def start_ex(message): """ Start command. Here we are starting state """ - bot.set_state(message.from_user.id, MyStates.name) + bot.set_state(message.from_user.id, MyStates.name, message.chat.id) bot.send_message(message.chat.id, 'Hi, write me a name') - +# Any state @bot.message_handler(state="*", commands='cancel') def any_state(message): """ Cancel state """ bot.send_message(message.chat.id, "Your state was cancelled.") - bot.delete_state(message.from_user.id) + bot.delete_state(message.from_user.id, message.chat.id) @bot.message_handler(state=MyStates.name) def name_get(message): @@ -36,8 +61,8 @@ def name_get(message): State 1. Will process when user's state is 1. """ bot.send_message(message.chat.id, f'Now write me a surname') - bot.set_state(message.from_user.id, MyStates.surname) - with bot.retrieve_data(message.from_user.id) as data: + bot.set_state(message.from_user.id, MyStates.surname, message.chat.id) + with bot.retrieve_data(message.from_user.id, message.chat.id) as data: data['name'] = message.text @@ -47,16 +72,16 @@ def ask_age(message): State 2. Will process when user's state is 2. """ bot.send_message(message.chat.id, "What is your age?") - bot.set_state(message.from_user.id, MyStates.age) - with bot.retrieve_data(message.from_user.id) as data: + bot.set_state(message.from_user.id, MyStates.age, message.chat.id) + with bot.retrieve_data(message.from_user.id, message.chat.id) as data: data['surname'] = message.text # result @bot.message_handler(state=MyStates.age, is_digit=True) def ready_for_answer(message): - with bot.retrieve_data(message.from_user.id) as data: + with bot.retrieve_data(message.from_user.id, 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") - bot.delete_state(message.from_user.id) + bot.delete_state(message.from_user.id, message.chat.id) #incorrect number @bot.message_handler(state=MyStates.age, is_digit=False) @@ -68,7 +93,4 @@ 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/async_telebot.py b/telebot/async_telebot.py index 0e297ee..a2edafb 100644 --- a/telebot/async_telebot.py +++ b/telebot/async_telebot.py @@ -364,12 +364,13 @@ class AsyncTeleBot: handler_error = None data = {} process_handler = True - middleware_result = await middleware.pre_process(message, data) - if isinstance(middleware_result, SkipHandler): - await middleware.post_process(message, data, handler_error) - process_handler = False - if isinstance(middleware_result, CancelUpdate): - return + if middleware: + middleware_result = await middleware.pre_process(message, data) + if isinstance(middleware_result, SkipHandler): + await middleware.post_process(message, data, handler_error) + process_handler = False + if isinstance(middleware_result, CancelUpdate): + return for handler in handlers: if not process_handler: break @@ -2481,8 +2482,8 @@ class AsyncTeleBot: """ return await asyncio_helper.delete_chat_photo(self.token, chat_id) - async def get_my_commands(self, scope: Optional[types.BotCommandScope]=None, - language_code: Optional[str]=None) -> List[types.BotCommand]: + async def get_my_commands(self, scope: Optional[types.BotCommandScope], + language_code: Optional[str]) -> List[types.BotCommand]: """ Use this method to get the current list of the bot's commands. Returns List of BotCommand on success. diff --git a/telebot/asyncio_handler_backends.py b/telebot/asyncio_handler_backends.py index 08db40f..2b7020c 100644 --- a/telebot/asyncio_handler_backends.py +++ b/telebot/asyncio_handler_backends.py @@ -17,3 +17,19 @@ class BaseMiddleware: async def post_process(self, message, data, exception): raise NotImplementedError + +class State: + def __init__(self) -> None: + self.name = None + def __str__(self) -> str: + return self.name + + +class StatesGroup: + def __init_subclass__(cls) -> None: + # print all variables of a subclass + for name, value in cls.__dict__.items(): + if not name.startswith('__') and not callable(value) and isinstance(value, State): + # change value of that variable + value.name = ':'.join((cls.__name__, name)) + diff --git a/telebot/handler_backends.py b/telebot/handler_backends.py index df4d37f..0b2bed7 100644 --- a/telebot/handler_backends.py +++ b/telebot/handler_backends.py @@ -148,4 +148,20 @@ class RedisHandlerBackend(HandlerBackend): self.clear_handlers(handler_group_id) return handlers + +class State: + def __init__(self) -> None: + self.name = None + def __str__(self) -> str: + return self.name + + +class StatesGroup: + def __init_subclass__(cls) -> None: + # print all variables of a subclass + for name, value in cls.__dict__.items(): + if not name.startswith('__') and not callable(value) and isinstance(value, State): + # change value of that variable + value.name = ':'.join((cls.__name__, name)) + \ No newline at end of file