Конечные автоматы (FSM)

Старая версия aiogram

Для aiogram 3.x эта глава расположена здесь.

Обновление главы (нажмите, чтобы показать)

Эта глава была обновлена в сентябре 2022 года, поскольку ранее содержала не самые лучшие советы по переходу между состояниями. Конкретно: не рекомендуется вызывать метод .set() у самого стейта; вместо этого лучше пользоваться методом set_state() у объекта FSMContext. Также не стоит пользоваться методом .next() по той же причине: неочевидное поведение.

Теория

В этой главе мы поговорим о, пожалуй, самой важной возможности ботов: о системе диалогов. К сожалению, далеко не все действия в боте можно выполнить за одно сообщение или команду. Предположим, есть бот для знакомств, где при регистрации нужно указать имя, возраст и отправить фотографию с лицом. Можно, конечно, попросить пользователя отправить фотографию, а в подписи к ней указать все данные, но это неудобно для обработки и запроса повторного ввода.
Теперь представим пошаговый ввод данных, где в начале бот «включает» режим ожидания определённой информации от конкретного юзера, далее на каждом этапе проверяет вводимые данные, а по команде /cancel прекращает ожидать очередной шаг и возвращается в основной режим. Взгляните на схему ниже:

Зелёной стрелкой обозначен процесс перехода по шагам без ошибок, синие стрелки означают сохранение текущего состояния и ожидание повторного ввода (например, если юзер указал, что ему 250 лет, следует запросить возраст заново), а красные показывают выход из всего процесса из-за команды /cancel или любой другой, означающей отмену.

Процесс со схемы выше в теории алгоритмов называется конечным автоматом (или FSM — Finite State Machine). Подробнее об этом можно прочесть здесь.

Практика

В aiogram механизм конечных автоматов реализован гораздо лучше, чем в том же pyTelegramBotAPI. Во фреймворк уже встроена поддержка различных бэкендов для хранения состояний между перезапусками бота (впрочем, никто не мешает написать свой), а помимо, собственно, состояний можно хранить произвольные данные, например, вышеописанные имя и возраст для последующего использования где-либо. Список имеющихся хранилищ FSM можно найти в репозитории aiogram, а в этой главе мы будем пользоваться самым простейшим бэкендом MemoryStorage, который хранит все данные в оперативной памяти. Он идеально подходит для примеров, но не рекомендуется использовать его в реальных проектах, т.к. MemoryStorage хранит все данные в оперативной памяти без сброса на диск. Также стоит отметить, что конечные автоматы можно использовать не только с обработчиками сообщений (message_handler, edited_message_handler), но также с колбэками и инлайн-режимом.

В качестве примера мы напишем имитатор заказа еды и напитков в кафе, а заодно научимся хранить логически разные хэндлеры в разных файлах.

Примечание об исходных текстах к главе

В тексте будет рассмотрен не весь код бота, некоторые импорты и обработчики специально пропущены для улучшения читабельности. Полный набор исходников можно найти на GitHub.

Благодарность

За основу структуры файлов и каталогов взят репозиторий tgbot_template от пользователя Tishka17. В этой главе будет рассмотрен сильно упрощённый вариант его примера, а далее по мере усложнения бота структура файлов будет расширяться.
Спасибо!

Структура файлов и каталогов

Точкой входа будет являться файл bot.py, рядом с ним каталог "config" с файлом конфигурации bot.ini. В прошлых главах была всего одна переменная, которую передавали через environment variables, но когда настроек становится много, хорошей идеей становится использование отдельного файла конфигурации и его парсинг стандартным питоновским Configparser-ом. Здесь же будет расположен пакет app, внутри которого расположен файл config.py, отвечающий за разбор файла конфигурации, а также пакет handlers с различными хэндлерами для логически разных наборов шагов. Схематично всё вышеперечисленное выглядит как-то так:

├── app/
│ ├── config.py
│ ├── handlers/
│ │ ├── common.py
│ │ ├── drinks.py
│ │ ├── food.py
│ │ └── __init__.py
│ └── __init__.py
├── config/
│ └── bot.ini
├── bot.py
└── requirements.txt

О модулях, пакетах и каталогах

Модули в Python — это файлы с расширением *.py. Для структурирования файлов их можно объединять в каталоги. Если в таком каталоге создать файл __init__.py, то каталог превращается в пакет. Подробнее обо всём этом (с примерами) можно прочесть на сайте devpractice.ru.

Создание шагов

Рассмотрим описание шагов для «заказа» еды. Для начала в файле app/handlers/food.py импортируем необходимые объекты и приведём списки блюд и их размеров (в реальной жизни эта информация может динамически подгружаться из какой-либо БД):

from aiogram import Dispatcher, types
from aiogram.dispatcher import FSMContext
from aiogram.dispatcher.filters.state import State, StatesGroup

# Эти значения далее будут подставляться в итоговый текст, отсюда 
# такая на первый взгляд странная форма прилагательных
available_food_names = ["суши", "спагетти", "хачапури"]
available_food_sizes = ["маленькую", "среднюю", "большую"]

Теперь опишем все возможные «состояния» конкретного процесса (выбор еды). На словах можно описать так: пользователь вызывает команду /food, бот отвечает сообщением с просьбой выбрать блюдо и встаёт в состояние *ожидает выбор блюда* для конкретного пользователя. Как только юзер делает выбор, бот, находясь в этом состоянии, проверяет корректность ввода, а затем принимает решение, запросить ввод повторно (без смены состояния) или перейти к следующему шагу *ожидает выбор размера порции*. Когда пользователь и здесь вводит корректные данные, бот отображает итоговый результат (содержимое заказа) и сбрасывает состояние. Позднее в этой главе мы научимся делать принудительный сброс состояния на любом этапе командой /cancel.

Итак, перейдём непосредственно к описанию состояний. Желательно их указывать именно в том порядке, в котором предполагается переход пользователя, это позволит немного упростить код. Для хранения состояний необходимо создать класс, наследующийся от класса StatesGroup, внутри него нужно создать переменные, присвоив им экземпляры класса State:

class OrderFood(StatesGroup):
    waiting_for_food_name = State()
    waiting_for_food_size = State()

Напишем обработчик первого шага, реагирующий на команду /food (регистрировать его будем позднее):

# Обратите внимание: есть второй аргумент
async def food_start(message: types.Message, state: FSMContext):
    keyboard = types.ReplyKeyboardMarkup(resize_keyboard=True)
    for name in available_food_names:
        keyboard.add(name)
    await message.answer("Выберите блюдо:", reply_markup=keyboard)
    await state.set_state(OrderFood.waiting_for_food_name.state)

В последней строке мы явно говорим боту встать в состояние waiting_for_food_name из группы OrderFood. Следующая функция будет вызываться только из указанного состояния, сохранять полученный от пользователя текст (если он валидный) и переходить к очередному шагу:

async def food_chosen(message: types.Message, state: FSMContext):
    if message.text.lower() not in available_food_names:
        await message.answer("Пожалуйста, выберите блюдо, используя клавиатуру ниже.")
        return
    await state.update_data(chosen_food=message.text.lower())

    keyboard = types.ReplyKeyboardMarkup(resize_keyboard=True)
    for size in available_food_sizes:
        keyboard.add(size)
    await state.set_state(OrderFood.waiting_for_food_size.state)
    await message.answer("Теперь выберите размер порции:", reply_markup=keyboard)

Разберём некоторые строки из блока выше отдельно. В определении функции food_chosen (строка 2) есть второй аргумент state типа FSMContext. Через него можно получить данные от FSM-бэкенда. В строке 2 производится проверка текста от пользователя. Если он ввёл произвольный текст, а не использовал кнопки, то необходимо сообщить об ошибке и досрочно завершить выполнение функции. При этом состояние пользователя останется тем же и бот снова будет ожидать выбор блюда.
К моменту перехода к строке 5 мы уже уверены, что пользователь указал корректное название блюда, поэтому можно спокойно сохранить полученный текст в хранилище данных FSM. Воспользуемся функцией update_data() и сохраним текст сообщения под ключом chosen_food и со значением message.text.lower().
В строке 10 мы готовы продвинуть пользователя на следующий шаг, и вызываем метод set_state(...) с нужным стейтом внутри.

Осталось реализовать последнюю функцию, которая отвечает за получение размера порции (с аналогичной проверкой ввода) и вывод результатов пользователю:

async def food_size_chosen(message: types.Message, state: FSMContext):
    if message.text.lower() not in available_food_sizes:
        await message.answer("Пожалуйста, выберите размер порции, используя клавиатуру ниже.")
        return
    user_data = await state.get_data()
    await message.answer(f"Вы заказали {message.text.lower()} порцию {user_data['chosen_food']}.\n"
                         f"Попробуйте теперь заказать напитки: /drinks", reply_markup=types.ReplyKeyboardRemove())
    await state.finish()

На что стоит обратить внимание: во-первых, получить хранимые данные из FSM можно методом get_data() у state. Во-вторых, т.к. это словарь, то и извлечение содержимого аналогично (user_data['chosen_food'], не забудьте, что в общем случае какого-то ключа может не быть, что приведёт к KeyError, пользуйтесь методом get()). В-третьих, вызов метода finish() сбрасывает не только состояние, но и хранящиеся данные. Если надо сбросить только состояние, воспользуйтесь await state.reset_state(with_data=False)

Наконец, напишем обычную функцию для регистрации вышестоящих обработчиков, на вход она будет принимать диспетчер. Значение состояния "*" при регистрации food_start() означает срабатывание при любом состоянии. Грубо говоря, если пользователь находится на шаге *выбор размера порции*, но решает начать заново и вводит /food, бот вызовет food_start() и начнёт весь процесс заново.

def register_handlers_food(dp: Dispatcher):
    dp.register_message_handler(food_start, commands="food", state="*")
    dp.register_message_handler(food_chosen, state=OrderFood.waiting_for_food_name)
    dp.register_message_handler(food_size_chosen, state=OrderFood.waiting_for_food_size)

Шаги для выбора напитков делаются совершенно аналогично. Попробуйте сделать самостоятельно или загляните в исходные тексты к этой главе.

Общие команды

Раз уж заговорили о сбросе состояний, давайте в файле app/handlers/common.py реализуем обработчики команды /start и действия «отмены». Первая должна показывать некий приветственный/справочный текст, а вторая просто пишет "действие отменено". Обе функции сбрасывают состояние и данные и убирают обычную клавиатуру, если вдруг она есть:

async def cmd_start(message: types.Message, state: FSMContext):
    await state.finish()
    await message.answer(
        "Выберите, что хотите заказать: напитки (/drinks) или блюда (/food).",
        reply_markup=types.ReplyKeyboardRemove()
    )

async def cmd_cancel(message: types.Message, state: FSMContext):
    await state.finish()
    await message.answer("Действие отменено", reply_markup=types.ReplyKeyboardRemove())

Зарегистрируем эти два обработчика:

def register_handlers_common(dp: Dispatcher):
    dp.register_message_handler(cmd_start, commands="start", state="*")
    dp.register_message_handler(cmd_cancel, commands="cancel", state="*")
    dp.register_message_handler(cmd_cancel, Text(equals="отмена", ignore_case=True), state="*")

Почему строк три, а не две? Дело в том, что один и тот же обработчик можно вызвать по разным событиям. Вот мы и зарегистрируем функцию cmd_cancel() для вызова как по команде /cancel, так и по отправке сообщения "Отмена" (в любом регистре). К слову, если вы навешиваете декораторы напрямую на функцию, то выглядеть это будет следующим образом:

@dp.message_handler(commands="cancel", state="*")
@dp.message_handler(Text(equals="отмена", ignore_case=True), state="*")
async def cmd_cancel(message: types.Message, state: FSMContext):
    ...

Точка входа

Вернёмся к файлу bot.py и реализуем две функции: set_commands() и main(). В первой зафиксируем список команд, доступных в интерфейсе Telegram при нажатии на кнопку Меню или [ / ], а во второй проведём необходимые действия по запуску бота: настроим логирование, загрузим и распарсим файл конфигурации, объявим объекты Bot и Dispatcher, зарегистрируем хэндлеры и команды и, наконец, запустим бота в режиме поллинга:

# Не забудьте про импорты

logger = logging.getLogger(__name__)

# Регистрация команд, отображаемых в интерфейсе Telegram
async def set_commands(bot: Bot):
    commands = [
        BotCommand(command="/drinks", description="Заказать напитки"),
        BotCommand(command="/food", description="Заказать блюда"),
        BotCommand(command="/cancel", description="Отменить текущее действие")
    ]
    await bot.set_my_commands(commands)


async def main():
    # Настройка логирования в stdout
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    )
    logger.error("Starting bot")

    # Парсинг файла конфигурации
    config = load_config("config/bot.ini")

    # Объявление и инициализация объектов бота и диспетчера
    bot = Bot(token=config.tg_bot.token)
    dp = Dispatcher(bot, storage=MemoryStorage())

    # Регистрация хэндлеров
    register_handlers_common(dp)
    register_handlers_drinks(dp)
    register_handlers_food(dp)

    # Установка команд бота
    await set_commands(bot)

    # Запуск поллинга
    # await dp.skip_updates()  # пропуск накопившихся апдейтов (необязательно)
    await dp.start_polling()

if __name__ == '__main__':
    asyncio.run(main())

Теперь, вооружившись знаниями о конечных автоматах, вы можете безбоязненно писать ботов с системой диалогов.