Как сделать telegram-бота для игры в Тайного Санту

Перед Новым годом мы организовали тайного санту. Для упрощения процесса задумались о боте. Да, мы нашли на просторах гитхаба различные варианты, но решили не лишать себя праздничного веселья от создания бота на коленке. Меня зовут Вильданов Ринат, я python разработчик в Технократии, и я расскажу, что мы наделали. Возможно, описание нашего пути поможет и вам.

Первым делом надо определиться с тем, что должен делать бот:

1. Регистрировать пользователя. Делаем это при старте бота

2. В компании мы хорошо общаемся между собой, но какие-то факты о потенциальном “внучке” могли ускользнуть, поэтому добавляем каждому пользователю биографию — по ней будет легче выбирать подарки

3. Не забываем про удаленщиков. Для них при заполнении биографии добавим стейт про адрес.

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

1. Добавим счетчик пользователей

2. Праздника не будет, если пары не распределены. Делаем команду для запуска события

3. Разделили создание пар и рассылку внучков на две разные команды

4. Админ говорит.  Добавили оповещалку пользователей от имени админа

5. Команда по очистке базы. Исключительно для тестовой части, чтобы не лезть каждый раз в БД.

6. Защитим права админа.  Люди у нас любознательные, кто-нибудь точно захочет побыть в роли админа-Гринча без прав на  это. Позаботимся об этом,  а виновникам добавим углей в список подарков. 

Определились с задачами, определимся со стеком. Самый простой вариант, как мне кажется,—  aiogram. Прекрасная документация, много проектов, в которых можно подглядеть реализацию. 

Надо где-то хранить данные. Для этого выберем postgres. Классика. Для взаимодействия с базой будем использовать sqlalchemy. Да, это мощно.

Чтобы комфортно накатывать миграции, добавим alembic. Кажется, все. Начнем разработку.

Расчехляем шаблон для бота и немного дорабатываем его под наши нужды. Имеем что то подобное:

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

telegram_users.py

class TelegramUser(Base):    __tablename__ = "telegram_user"    id = Column(Integer, primary_key=True)    first_name = Column(String)    last_name = Column(String)    tg_id = Column(Integer, unique=True)    chat_id = Column(Integer, unique=True)    description = Column(String, nullable=True)    address = Column(String, nullable=True) 

event.py

class Event(Base):    __tablename__ = "event"    id = Column(Integer, primary_key=True)    user_id = Column(Integer, ForeignKey("telegram_user.id"), unique=True)  # Кому дарит    santa_id = Column(Integer, ForeignKey("telegram_user.id"), unique=True)  # Кто дарит 

Ничего специфического здесь нет: описываем модели, на основе которых будут сгенерированы миграции

Классы готовы, но сам alembic их не найдет, ему необходимо помочь. Для этого перейдем в alembic/env.py и объявим, на основе чего делать миграции.

Теперь надо зафиксировать, куда катить миграции.  Для этого в def run_migrations_offline укажем DSN. DSN лежит в файле settings.py со всеми настройками.

def run_migrations_offline():    url = DSN     context.configure(        url=url,        target_metadata=target_metadata,        literal_binds=True,        dialect_opts={"paramstyle": "named"},    )     with context.begin_transaction():        context.run_migrations() 

Так же поправим def run_migrations_online

def run_migrations_online()    conf = config.get_section(config.config_ini_section)    conf["sqlalchemy.url"] = SYNC_DSN    connectable = engine_from_config(        conf,        prefix="sqlalchemy.",        poolclass=pool.NullPool,    )     with connectable.connect() as connection:        context.configure(connection=connection, target_metadata=target_metadata)         with context.begin_transaction():            context.run_migrations() 

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

Генерируем миграции и накатываем:

Создание: alembic revision --autogenerate  Накатывание: alembic upgrade head

Замечательно, таблицы создались, займемся реализациями комманд. Для этого пройдем в модуль handlers.

Разбирать реализацию всех команд не вижу смысла: если понятно, как сделать хотя бы одну, с остальными проблем не возникнет.

@dp.message_handler(commands=["start"]) async def start(message: types.Message):    data = {        "chat_id": message["chat"]["id"],        "tg_id": message["from"]["id"],        "first_name": message["chat"]["first_name"],        "last_name": message["chat"]["last_name"],    }    response = await TelegramUser.add_user(data)    await TelegramUser.get_user_by_id(int(message["from"]["id"]))    await message.answer(response)

Из переменной message с соответствующим типом вытаскиваем необходимые для нас данные: id пользователя, id чата (они совпадают, но раз подаются отдельно, то и хранить мы их будем отдельно), также имя и фамилию. Передаем эти данные в метод add_user класса, на основе которого была создана таблица. Внутри данного метода будет простое добавление данных в базу.

Думаю, с пользовательским ручками разобрались. С админскими +- тоже, за одним исключением — необходимо реализовать проверку на админа. Первое, что приходит в голову, — решение в лоб: в каждом админском методе сравнивать id владельца и id пользователя, который отправил сообщение. Метод имеет право на существование, так как решает поставленную задачу, но он не идеален, так как будет много дублирующего кода.

Можно написать декоратор и обернуть все админские ручки в него. И да и нет. Декоратор мы использовать будем, но писать будем не его, а фильтр. Обращаемся к документации и видим там параграф про фильтры. Это то, что нам нужно, копипастим кусок кода, переименовываем и получаем подобное:

class AdminFilter(BoundFilter):    key = 'is_admin'     def __init__(self, is_admin):        self.is_admin = is_admin     async def check(self, message: types.Message):        if message["from"]["id"] not in OWNER_ID:            return False        return True 

Как можно заметить, проверка происходит в методе check.

Теперь к  расстановке пар. Тут нет единого способа решения. К примеру, вот свежая статья про это. Мы решили не идти против течения и воспользовались обычным шаффлом из модуля random. 

Получаем список всех id. Смешиваем их в случайном порядке и получаем такие пары: пользователь c индексом 0 дарит подарок пользователю с индексом 1. Пользователь c индексом 1 дарит пользователю с индексом 2 … пользователь c индексом (n-2) дарит пользователю с индексом (n-1), пользователь с индексом (n-1) дарит подарок пользователь с индексом 0. Говоря про индексы, имеется в виду индекс списка, а значением является id. Вот собственно и вся магия.

Пары получены, теперь осталось рассказать сантам про их подопечных. Для этого мы запланировали команду /notify

@dp.message_handler(commands=["notify"]) async def start_event(message: types.Message):    tg_id = message["from"]["id"]    responses = await Event.notify_all(int(tg_id))    if isinstance(responses, str):        await message.answer(responses)        return    for response in responses:        await bot.send_message(response["santa_id"], format_message(response)) 

Тут тоже  несложно. С помощью таблицы event получаем список сант, матчим список подопечных с их адресами и описанием. Потом идем в цикле и с помощью telegram id отправляем сообщения.

Выводы

Бота запустили, подопечных всем раздали, успели получить первые отзывы и вот что мы не учли: 

• многие нажали на /start просто из интереса, а их уже зарегистрировало — нужно сделать подтверждение участия 

• у части ребят нет фамилии или она скрыта, из-за этого приходилось вручную искать какой из трех Алмазов кому выпал — надо сохранять alias

• хорошо бы сделать напоминание по таймеру для тех, кто не заполнил описание 

• поскольку все писалось на коленке, проверок маловато, хорошо бы еще понимать, кому сообщение отправилось, а кому нет

А вы играете в «Тайного Санту» на работе или с друзьями?

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