Перед Новым годом мы организовали тайного санту. Для упрощения процесса задумались о боте. Да, мы нашли на просторах гитхаба различные варианты, но решили не лишать себя праздничного веселья от создания бота на коленке. Меня зовут Вильданов Ринат, я 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

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

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

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

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