Перед Новым годом мы организовали тайного санту. Для упрощения процесса задумались о боте. Да, мы нашли на просторах гитхаба различные варианты, но решили не лишать себя праздничного веселья от создания бота на коленке. Меня зовут Вильданов Ринат, я python разработчик в Технократии, и я расскажу, что мы наделали. Возможно, описание нашего пути поможет и вам.
Первым делом надо определиться с тем, что должен делать бот:
Регистрировать пользователя. Делаем это при старте бота
В компании мы хорошо общаемся между собой, но какие-то факты о потенциальном “внучке” могли ускользнуть, поэтому добавляем каждому пользователю биографию — по ней будет легче выбирать подарки
Не забываем про удаленщиков. Для них при заполнении биографии добавим стейт про адрес.
С командами для обычного пользователя все. Теперь об админах.
Добавим счетчик пользователей
Праздника не будет, если пары не распределены. Делаем команду для запуска события
Разделили создание пар и рассылку внучков на две разные команды
Админ говорит. Добавили оповещалку пользователей от имени админа
Команда по очистке базы. Исключительно для тестовой части, чтобы не лезть каждый раз в БД.
Защитим права админа. Люди у нас любознательные, кто-нибудь точно захочет побыть в роли админа-Гринча без прав на это. Позаботимся об этом, а виновникам добавим углей в список подарков.
Определились с задачами, определимся со стеком. Самый простой вариант, как мне кажется,— 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
хорошо бы сделать напоминание по таймеру для тех, кто не заполнил описание
поскольку все писалось на коленке, проверок маловато, хорошо бы еще понимать, кому сообщение отправилось, а кому нет
А вы играете в «Тайного Санту» на работе или с друзьями?
Также подписывайтесь на наш телеграм-канал «Голос Технократии». Каждое утро мы публикуем новостной дайджест из мира ИТ, а по вечерам делимся интересными и полезными мастридами.