Урок 3. Машина состояний и то самое логгирование

Важно! Не забывайте обновлять библиотеку командой python3.6 -m pip install -U aiogram, так как разработчик постоянно обновляет её, избавляя от различных багов. Например в последнем релизе исправлены ошибки Fatal Python error: PyImport_GetModuleDict: no module dictionary! и передача списков в фильтр состояний.

Урок проводится с использованием aiogram версии 1.2

Сегодня мы научимся использовать:

• машину состояний

• встроенный middleware для логгирования

• Helper

Традиционно код урока доступен на GitHub

Создаем состояния

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

Итак, запишем в файл utils.py наш демонстрационный класс с состояниями:

from aiogram.utils.helper import Helper, HelperMode, ListItem


class TestStates(Helper):
    mode = HelperMode.snake_case

    TEST_STATE_0 = ListItem()
    TEST_STATE_1 = ListItem()
    TEST_STATE_2 = ListItem()
    TEST_STATE_3 = ListItem()
    TEST_STATE_4 = ListItem()
    TEST_STATE_5 = ListItem()

Обращу внимание читателя на то, что здесь нам интересно использовать именно ListItem, а не Item, так как в таком случае мы сможем использовать сложение разных состояний для передачи в handler (об этом дальше).

Если вы ещё не знакомы с классом Helper, то советую позже ознакомиться, а сейчас можно просто посмотреть на значения всех элементов (или каждого по отдельности, если захотите):

print(TestStates.all())

# ['test_state_0', 'test_state_1', 'test_state_2', 'test_state_3', 'test_state_4', 'test_state_5']

Так как в этот раз нам придется отправлять много сообщений, для разнообразия вынесем их в messages.py - для красоты.

Ещё не забываем добавить в config.py токен своего бота и мы готовы писать логику!

Указываем хранилище состояний и включаем логгирование

К привычным с прошлых уроков импортам у нас добавляется ещё парочка, а именно:

from aiogram.contrib.fsm_storage.memory import MemoryStorage
from aiogram.contrib.middlewares.logging import LoggingMiddleware

И тут же применяем их:

dp = Dispatcher(bot, storage=MemoryStorage())
dp.middleware.setup(LoggingMiddleware())

На первой строчке мы указали хранилище состояний в оперативной памяти, так как потеря этих состояний нам не страшна (да и этот вариант больше всего подходит для демонстрационных целей, так как не требует настройки). Однако если у вас от состояний что-то зависит, рекомендуется ипользовать более надеждное хранилище. На данный момент можно подключить Redis и RethinkDB.

На второй строчке подключаем логгирование. На нем долго останавливаться не буду - лучше проверьте его работу самостоятельно: например, можно добавить хэндлер только на текстовые сообщения, тогда при отправке стикера, фото (и т.п.), поста в канал, где бот является администратором, в логах увидите, что эти апдейты не отработаны никаким хэндлером.

Обрабатываем входящие сообщения

Итак, по традиции добавляем обработчики команд start и help:

@dp.message_handler(commands=['start'])
async def process_start_command(message: types.Message):
    await message.reply(MESSAGES['start'])


@dp.message_handler(commands=['help'])
async def process_help_command(message: types.Message):
    await message.reply(MESSAGES['help'])

А так же "ловим" все сообщения, отправленные при "нулевом" состоянии:

@dp.message_handler()
async def echo_message(msg: types.Message):
    await bot.send_message(msg.from_user.id, msg.text)

Переходим к главной теме нашего урока: состояниям

Эти самые состояния нужно как-то устанавливать, поэтому сделаем так:

@dp.message_handler(state='*', commands=['setstate'])
async def process_setstate_command(message: types.Message):
    argument = message.get_args()
    state = dp.current_state(user=message.from_user.id)
    if not argument:
        await state.reset_state()
        return await message.reply(MESSAGES['state_reset'])

    if (not argument.isdigit()) or (not int(argument) < len(TestStates.all())):
        return await message.reply(MESSAGES['invalid_key'].format(key=argument))

    await state.set_state(TestStates.all()[int(argument)])
    await message.reply(MESSAGES['state_change'], reply=False)

Не забываем, что хэндлеры обрабатываются в порядке их расположения в коде, поэтому описанная выше функция должна идти раньше приведенных ниже обработчиков.

В функции process_setstate_command мы проверяем, идут ли с командой какие-то аргументы и запрашиваем текущее состояние пользователя. Если никаких аргументов с командой не идёт, сбрасываем состояние. Если же аргументы есть, то проверяем, чтобы те соответствовали нашему условию: не отрицательное число (минус в строке не пройдет валидацию .isdigit()), которое меньше количества всех состояний. Если аргумент не подходит, сообщаем пользователю об этом. В ином случае устанавливаем новое текущее состояние и даем фидбек на действие сообщением. Обращу внимание читателя также на конструкцию message.reply('Some text', reply=False). Указав reply=False мы используем шорткат, позволяющий нам ответить в тот же чат, не доставая айди пользователя / чата как, например, вот тут.

Теперь отрабатываем входящие сообщения при выбранном состоянии

@dp.message_handler(state=TestStates.TEST_STATE_1)
async def first_test_state_case_met(message: types.Message):
    await message.reply('Первый!', reply=False)

В данном хэндлере мы передаем в состояние TestStates.TEST_STATE_1, что эквивалентно ['test_state_1'] - отмечу дважды, тут передается именно массив.

Теперь добавим такой хэндлер:

@dp.message_handler(state=TestStates.TEST_STATE_2[0])
async def second_test_state_case_met(message: types.Message):
    await message.reply('Второй!', reply=False)

Здесь мы уже указываем состояние 'test_state_2'.

Библиотека сама понимает, когда мы передаем список состояний, а когда только одно состояние и под капотом обрабатывает их по-разному, но нам нет смысла об этом задумываться. В этом плане мы в плюсе, так как можем сделать вот так:

@dp.message_handler(state=TestStates.TEST_STATE_3 | TestStates.TEST_STATE_4)
async def third_or_fourth_test_state_case_met(message: types.Message):
    await message.reply('Третий или четвертый!', reply=False)

Выражение TestStates.TEST_STATE_3 | TestStates.TEST_STATE_4 возвращает массив: ['test_state_3', 'test_state_4']. Грубо говоря, при проверке состояний библиотека проходится по списку и если встречает совпадение, хэндлер отрабатывается.

Ну и последний на сегодня хэндлер:

@dp.message_handler(state=TestStates.all())
async def some_test_state_case_met(message: types.Message):
    with dp.current_state(user=message.from_user.id) as state:
        text = MESSAGES['current_state'].format(
            current_state=await state.get_state(),
            states=TestStates.all()
        )
    await message.reply(text, reply=False)

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

Ещё в приведённом выше блоке кода используется контекстный менеджер (with as) - многим с ним удобнее.

Для красоты ещё стоит закрывать соединение с хранилищем состояний, для этого объявляем функцию:

async def shutdown(dispatcher: Dispatcher):
    await dispatcher.storage.close()
    await dispatcher.storage.wait_closed()

И в поллере указываем её:

if __name__ == '__main__':
    executor.start_polling(dp, on_shutdown=shutdown)

Вот и всё!