From f972677cb18160be0a85e7622aaf3a6840e543ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D0=BD=D0=B4=D1=80=D0=B5=D0=B9=20=D0=9C=D0=B8=D1=88?= =?UTF-8?q?=D1=83=D1=80=D0=B8=D0=BD?= Date: Thu, 4 Dec 2025 03:39:06 +0300 Subject: [PATCH 1/4] refactor: general refactoring --- .env.example | 4 +- random_prizes.py | 17 ++--- src/bot.py | 181 ++++++++++++++--------------------------------- src/callbacks.py | 26 +++++++ src/models.py | 15 ++-- src/templates.py | 59 +++++++++++++++ src/userdb.py | 98 +++++++------------------ 7 files changed, 180 insertions(+), 220 deletions(-) create mode 100644 src/callbacks.py create mode 100644 src/templates.py diff --git a/.env.example b/.env.example index 00adfda..75680fc 100644 --- a/.env.example +++ b/.env.example @@ -1,4 +1,6 @@ TG_BOT_TOKEN="" MONGODB_USERNAME="" MONGODB_PASSWORD="" -MONGODB_PORT=27017 \ No newline at end of file +MONGODB_PORT=27017 + +MONGODB_HOST="mongodb://login:password@ip:port" # If you don't use docker compose \ No newline at end of file diff --git a/random_prizes.py b/random_prizes.py index e10f0b6..0013c99 100644 --- a/random_prizes.py +++ b/random_prizes.py @@ -1,24 +1,19 @@ -from dotenv import load_dotenv -import os import argparse -import pymongo +import os import random +import pymongo +from dotenv import load_dotenv + load_dotenv() -MONGODB_USERNAME = os.getenv("MONGODB_USERNAME") -MONGODB_PASSWORD = os.getenv("MONGODB_PASSWORD") -MONGODB_PORT = os.getenv("MONGODB_PORT") -MONGODB_IP = os.getenv("MONGODB_IP") +MONGODB_HOST = os.getenv("MONGODB_HOST") parser = argparse.ArgumentParser() parser.add_argument("count", type=int) args = parser.parse_args() -MONGO_HOST = ( - f"mongodb://{MONGODB_USERNAME}:{MONGODB_PASSWORD}@{MONGODB_IP}:{MONGODB_PORT}" -) -client = pymongo.MongoClient(MONGO_HOST) +client = pymongo.MongoClient(MONGODB_HOST) database = client.get_database("cu_graph_bot") collection = database.get_collection("users") users = ["@" + i["username"] for i in collection.find({"_links.4": {"$exists": True}})] diff --git a/src/bot.py b/src/bot.py index 03d5c9e..4f6d5cf 100644 --- a/src/bot.py +++ b/src/bot.py @@ -1,23 +1,22 @@ import os -from typing import List -from aiogram import Bot, Dispatcher, F, html, types +from aiogram import Bot, Dispatcher, F, types +from aiogram.client.default import DefaultBotProperties from aiogram.enums import ParseMode -from aiogram.filters import Command, CommandStart, callback_data +from aiogram.filters import CommandStart from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup from aiogram.types import ( - BotCommand, CallbackQuery, InlineKeyboardButton, InlineKeyboardMarkup, KeyboardButton, ReplyKeyboardMarkup, - user, ) -from aiogram.utils.formatting import Bold, CustomEmoji, Text +from aiogram.utils.formatting import as_list from dotenv import load_dotenv +from . import callbacks, templates from .models import Link, User, check_tg_username from .userdb import userdb @@ -27,31 +26,10 @@ if TOKEN is None: raise Exception("Couldn't find TG_BOT_TOKEN") -bot = Bot(token=TOKEN) +bot = Bot(token=TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML)) dp = Dispatcher() -class StartingCallback(callback_data.CallbackData, prefix="start"): - pass - - -class LinkCallback(callback_data.CallbackData, prefix="link"): - username_to: str - rating: int - - -class SexCallback(callback_data.CallbackData, prefix="sex"): - sex: str - - -class CourseCallback(callback_data.CallbackData, prefix="course"): - course: int - - -class LivingCallback(callback_data.CallbackData, prefix="living"): - living: str - - rkb = ReplyKeyboardMarkup( keyboard=[ [KeyboardButton(text="Мои контакты")], @@ -60,50 +38,6 @@ class LivingCallback(callback_data.CallbackData, prefix="living"): ] ) -starting_message = """🧬 ДОБРО ПОЖАЛОВАТЬ В CAMPUS DNA! - -Мы создаём первую карту социальных связей нашего университета. -Это исследование научной студии, и каждый участник получит: -• Личный анализ социального типа -• Место на интерактивной карте универа -• Шанс выиграть КРУТЫЕ ПРИЗЫ 🎁 - -🏆 УСЛОВИЯ УЧАСТИЯ В РОЗЫГРЫШЕ: - -1. ✅ Подписаться на канал @campusdna -2. ✅ Отметь 5+ друзей и оцени вашу близость - -🎁 ПРИЗОВОЙ ФОНД: -• 20 ПИЦЦ (1 пицца = 1 победитель) -• ИГРУШКИ-МИНЬОНЫ -• МЕРЧ ОТ ЦУ И ПАРТНЁРОВ - -📢 Следи за розыгрышами в канале: @campusdna - -🧭 ЧТО ДЕЛАТЬ ДАЛЬШЕ: - -Сначала тебе нужно ввести базовые сведения: пол, курс, общежитие. -Затем я попрошу тебя ввести юзернеймы твоих друзей в Telegram -и оценить вашу близость по шкале от 1 до 3. - -Чем больше друзей ты отметишь — тем точнее будет твой -социальный портрет и тем ценнее твой вклад в исследование! - -Готов начать и узнать, кто ты в социальной сети университета?""" - -explaining_links_message = """⁉️ НА КАКИЕ ГРУППЫ МЫ ДЕЛИМ СВЯЗИ? -🔴 1 — Друзья -«С ними я провожу больше всего времени» -Постоянное общение в универе и в мессенджерах. Видимся почти каждый день. Делимся личными новостями, поддерживаем друг друга. - -🟡 2 — Приятели -«Всегда подойду спросить: "Как дела? Как жизнь?"» -Видимся несколько раз в неделю. Общаемся и про учебу, и про жизнь, иногда затрагиваем личное (но не глубокое). Можем вместе пообедать или поиграть в пин-понг. - -🔵 3 — Знакомые -«Мы здороваемся в коридоре» -Видимся изредка, общение короткое и ситуативное. В основном на учебные/повседневные темы.""" - class AddingUser(StatesGroup): sex = State() @@ -114,26 +48,25 @@ class AddingUser(StatesGroup): @dp.message(CommandStart()) async def start_handler(message: types.Message, state: FSMContext): await message.answer( - starting_message, + templates.starting_message, reply_markup=InlineKeyboardMarkup( inline_keyboard=[ [ InlineKeyboardButton( - text="Далее", callback_data=StartingCallback().pack() + text="Далее", callback_data=callbacks.StartingCallback().pack() ) ] ] ), - parse_mode=ParseMode.HTML, ) -@dp.callback_query(StartingCallback.filter()) +@dp.callback_query(callbacks.StartingCallback.filter()) async def next_handler( - query: CallbackQuery, callback_data: StartingCallback, state: FSMContext + query: CallbackQuery, callback_data: callbacks.StartingCallback, state: FSMContext ): if await userdb.get_user(query.from_user.username) is not None: - await start_survey(query.message) + await explaining_links(query.message) return await state.set_state(AddingUser.sex) await question_sex(query.message, state) @@ -148,10 +81,12 @@ async def question_sex(message: types.Message, state: FSMContext): inline_keyboard=[ [ InlineKeyboardButton( - text="Мужской", callback_data=SexCallback(sex="male").pack() + text="Мужской", + callback_data=callbacks.SexCallback(sex="male").pack(), ), InlineKeyboardButton( - text="Женский", callback_data=SexCallback(sex="female").pack() + text="Женский", + callback_data=callbacks.SexCallback(sex="female").pack(), ), ] ] @@ -159,9 +94,9 @@ async def question_sex(message: types.Message, state: FSMContext): ) -@dp.callback_query(SexCallback.filter()) +@dp.callback_query(callbacks.SexCallback.filter()) async def process_sex( - query: CallbackQuery, callback_data: SexCallback, state: FSMContext + query: CallbackQuery, callback_data: callbacks.SexCallback, state: FSMContext ): await state.update_data(sex=callback_data.sex) await question_course(query.message, state) @@ -176,10 +111,12 @@ async def question_course(message: types.Message, state: FSMContext): inline_keyboard=[ [ InlineKeyboardButton( - text="1", callback_data=CourseCallback(course=1).pack() + text="1", + callback_data=callbacks.CourseCallback(course=1).pack(), ), InlineKeyboardButton( - text="2", callback_data=CourseCallback(course=2).pack() + text="2", + callback_data=callbacks.CourseCallback(course=2).pack(), ), ] ] @@ -187,9 +124,9 @@ async def question_course(message: types.Message, state: FSMContext): ) -@dp.callback_query(CourseCallback.filter()) +@dp.callback_query(callbacks.CourseCallback.filter()) async def process_course( - query: CallbackQuery, callback_data: CourseCallback, state: FSMContext + query: CallbackQuery, callback_data: callbacks.CourseCallback, state: FSMContext ): await state.update_data(course=callback_data.course) await question_living(query.message, state) @@ -205,25 +142,27 @@ async def question_living(message: types.Message, state: FSMContext): [ InlineKeyboardButton( text="В Облаке", - callback_data=LivingCallback(living="Cloud").pack(), + callback_data=callbacks.LivingCallback(living="Cloud").pack(), ), ], [ InlineKeyboardButton( text="В Космосе", - callback_data=LivingCallback(living="Cosmos").pack(), + callback_data=callbacks.LivingCallback(living="Cosmos").pack(), ), ], [ InlineKeyboardButton( text="В Байкале", - callback_data=LivingCallback(living="Baikal").pack(), + callback_data=callbacks.LivingCallback(living="Baikal").pack(), ), ], [ InlineKeyboardButton( text="Не в общаге", - callback_data=LivingCallback(living="Homeless").pack(), + callback_data=callbacks.LivingCallback( + living="Homeless" + ).pack(), ), ], ] @@ -231,9 +170,9 @@ async def question_living(message: types.Message, state: FSMContext): ) -@dp.callback_query(LivingCallback.filter()) +@dp.callback_query(callbacks.LivingCallback.filter()) async def process_living( - query: CallbackQuery, callback_data: LivingCallback, state: FSMContext + query: CallbackQuery, callback_data: callbacks.LivingCallback, state: FSMContext ): state_data = await state.get_data() await userdb.add_user( @@ -244,12 +183,19 @@ async def process_living( living=callback_data.living, ) ) - await start_survey(query.message) + await explaining_links(query.message) -async def start_survey(message: types.Message): - await message.answer(explaining_links_message, parse_mode=ParseMode.HTML) +async def explaining_links(message: types.Message): await message.answer( + templates.explaining_links_message, + callback_data=callbacks.TypeInfoCallback().pack(), + ) + + +@dp.callback_query(callbacks.TypeInfoCallback) +async def start_survey(query: CallbackQuery, callback_data: callbacks.TypeInfoCallback): + await query.answer( "Напиши юзернейм (@username) и я предложу тебе выбрать его категорию", reply_markup=rkb, ) @@ -277,20 +223,6 @@ async def get_usS(message: types.Message): await message.answer(all_users_and_rating, reply_markup=rkb) -def make_type_str(type: str, profile: str, strong_sides: List[str], recomendation: str): - return f"""🎯ТИП: «{type}» - -📊 Ваш профиль: -{profile} - -💪 Ваши сильные стороны: -• {"\n• ".join(strong_sides)} - -🌟 Рекомендация: -{recomendation} -""" - - @dp.message(F.text == "Узнать тип личности") async def get_summary(message: types.Message): links = await userdb.get_links(message.from_user.username) @@ -304,7 +236,7 @@ async def get_summary(message: types.Message): ) elif p3 >= 0.5: await message.answer( - make_type_str( + templates.make_type_str( "Сердце компании", "Вы создаете глубокие, осознанные отношения. Для вас важно не количество контактов, а их качество и надежность", [ @@ -315,11 +247,10 @@ async def get_summary(message: types.Message): ], 'Попробуйте иногда быть "социальным мостом" — знакомить своих друзей из разных кругов. Ваша глубина общения может стать основой для новых интересных компаний', ), - parse_mode=ParseMode.HTML, ) elif p2 >= 0.4: await message.answer( - make_type_str( + templates.make_type_str( "Социальный организатор", "Вы — мастер поддерживать ровные, комфортные отношения. С вами легко и приятно общаться на повседневные темы", [ @@ -330,11 +261,10 @@ async def get_summary(message: types.Message): ], "Попробуйте выбрать 1-2 самых интересных вам приятеля и предложить им более тесное общение — совместный проект или регулярные встречи. Ваши легкие связи могут перерасти в нечто большее", ), - parse_mode=ParseMode.HTML, ) elif p3 >= 0.25 and p2 >= 0.25 and p1 >= 0.25: await message.answer( - make_type_str( + templates.make_type_str( "Универсальный коннектор", "Вы легко перемещаетесь между разными социальными слоями. От тактических знакомств до близкой дружбы — вы чувствуете себя комфортно на любом уровне", [ @@ -345,11 +275,10 @@ async def get_summary(message: types.Message): ], "Используйте свой дар соединять людей! Организуйте мини-встречи людей из разных ваших кругов — возможно, вы создадите новые интересные коллаборации", ), - parse_mode=ParseMode.HTML, ) elif p3 >= 0.35 and p1 >= 0.25: await message.answer( - make_type_str( + templates.make_type_str( "Стратегический коммуникатор", "Вы сочетаете глубокую эмоциональную привязанность с широким кругом полезных контактов. Это редкий и ценный навык!", [ @@ -360,11 +289,10 @@ async def get_summary(message: types.Message): ], 'Подумайте, как ваши "знакомые" могут помочь вашим "друзьям" (и наоборот). Вы идеально positioned для создания синергии между разными частями вашей сети', ), - parse_mode=ParseMode.HTML, ) elif abs(p3 - p2) <= 0.3 and abs(p2 - p1) / len(ratings) <= 0.3: await message.answer( - make_type_str( + templates.make_type_str( "Стабильный якорь", "Вы выстраиваете гармоничную социальную экосистему, где каждому типу отношений находится свое место", [ @@ -375,7 +303,6 @@ async def get_summary(message: types.Message): ], 'Ваша сила — в стабильности. Подумайте, не хотите ли вы немного "сдвинуть баланс" в какую-то сторону: углубить несколько связей или, наоборот, расширить круг тактических контактов', ), - parse_mode=ParseMode.HTML, ) else: await message.answer( @@ -405,7 +332,7 @@ async def user_name_checker(message: types.Message): [ InlineKeyboardButton( text="Близкий друг", - callback_data=LinkCallback( + callback_data=callbacks.LinkCallback( username_to=username_to, rating=3 ).pack(), ) @@ -413,7 +340,7 @@ async def user_name_checker(message: types.Message): [ InlineKeyboardButton( text="Приятель", - callback_data=LinkCallback( + callback_data=callbacks.LinkCallback( username_to=username_to, rating=2 ).pack(), ), @@ -421,7 +348,7 @@ async def user_name_checker(message: types.Message): [ InlineKeyboardButton( text="Знакомый", - callback_data=LinkCallback( + callback_data=callbacks.LinkCallback( username_to=username_to, rating=1 ).pack(), ), @@ -432,11 +359,9 @@ async def user_name_checker(message: types.Message): await message.answer("Кто он для тебя?", reply_markup=kb) -@dp.callback_query(LinkCallback.filter()) -async def process_data(query: CallbackQuery, callback_data: LinkCallback): +@dp.callback_query(callbacks.LinkCallback.filter()) +async def process_data(query: CallbackQuery, callback_data: callbacks.LinkCallback): from_username = query.from_user.username - if from_username is None: - raise Exception("WTF?") if await userdb.get_user(from_username) is None: await query.answer("Похоже тебе нужно перезапустить бота: /start") return @@ -445,10 +370,10 @@ async def process_data(query: CallbackQuery, callback_data: LinkCallback): Link(username_to=callback_data.username_to, rating=callback_data.rating), ) await query.message.edit_text( - **Text( + as_list( f"✅ @{callback_data.username_to} добавлен как {rating_to_text(callback_data.rating).lower()}", "\n📝 Чтобы добавить ещё друга — просто введи следующий юзернейм.", "\n🔁 Чем больше друзей ты добавишь — тем точнее будет твой социальный портрет!", - ).as_kwargs(), + ), reply_markup=None, ) diff --git a/src/callbacks.py b/src/callbacks.py new file mode 100644 index 0000000..b978398 --- /dev/null +++ b/src/callbacks.py @@ -0,0 +1,26 @@ +from aiogram.filters import callback_data + + +class StartingCallback(callback_data.CallbackData, prefix="start"): + pass + + +class LinkCallback(callback_data.CallbackData, prefix="link"): + username_to: str + rating: int + + +class SexCallback(callback_data.CallbackData, prefix="sex"): + sex: str + + +class CourseCallback(callback_data.CallbackData, prefix="course"): + course: int + + +class LivingCallback(callback_data.CallbackData, prefix="living"): + living: str + + +class TypeInfoCallback(callback_data.CallbackData, prefix="typeinfo"): + pass diff --git a/src/models.py b/src/models.py index a6fcfe5..985e847 100644 --- a/src/models.py +++ b/src/models.py @@ -5,25 +5,28 @@ def check_tg_username(username: str) -> str: - username = username.strip().strip('@') + username = username.strip().strip("@") pattern = r"^[A-z0-9_]+$" if re.match(pattern, username): return username - raise ValueError(f'{username} is not valid username') + raise ValueError(f"{username} is not valid username") + Username = Annotated[str, AfterValidator(check_tg_username)] Sex = Literal["male", "female"] Living = Literal["Cloud", "Cosmos", "Baikal", "Homeless"] + class Link(BaseModel): username_to: Username - rating: int = Field(ge=1,le=3) + rating: int = Field(ge=1, le=3) + -class User(BaseModel) : +class User(BaseModel): username: Username # TODO: Add sex, course, etc. sex: Sex - course: int = Field(ge=1,le=2) + course: int = Field(ge=1, le=2) living: Living _links: list[Link] = [] @@ -32,4 +35,4 @@ def set_link(self, link: Link): if i.username_to == link.username_to: i = link return - self._links.append(link) \ No newline at end of file + self._links.append(link) diff --git a/src/templates.py b/src/templates.py new file mode 100644 index 0000000..a6ac6f9 --- /dev/null +++ b/src/templates.py @@ -0,0 +1,59 @@ +from typing import List + +starting_message = """🧬 ДОБРО ПОЖАЛОВАТЬ В CAMPUS DNA! + +Мы создаём первую карту социальных связей нашего университета. +Это исследование научной студии, и каждый участник получит: +• Личный анализ социального типа +• Место на интерактивной карте универа +• Шанс выиграть КРУТЫЕ ПРИЗЫ 🎁 + +🏆 УСЛОВИЯ УЧАСТИЯ В РОЗЫГРЫШЕ: + +1. ✅ Подписаться на канал @campusdna +2. ✅ Отметь 5+ друзей и оцени вашу близость + +🎁 ПРИЗОВОЙ ФОНД: +• 20 ПИЦЦ (1 пицца = 1 победитель) +• ИГРУШКИ-МИНЬОНЫ +• МЕРЧ ОТ ЦУ И ПАРТНЁРОВ + +📢 Следи за розыгрышами в канале: @campusdna + +🧭 ЧТО ДЕЛАТЬ ДАЛЬШЕ: + +Сначала тебе нужно ввести базовые сведения: пол, курс, общежитие. +Затем я попрошу тебя ввести юзернеймы твоих друзей в Telegram +и оценить вашу близость по шкале от 1 до 3. + +Чем больше друзей ты отметишь — тем точнее будет твой +социальный портрет и тем ценнее твой вклад в исследование! + +Готов начать и узнать, кто ты в социальной сети университета?""" + +explaining_links_message = """⁉️ НА КАКИЕ ГРУППЫ МЫ ДЕЛИМ СВЯЗИ? +🔴 1 — Друзья +«С ними я провожу больше всего времени» +Постоянное общение в универе и в мессенджерах. Видимся почти каждый день. Делимся личными новостями, поддерживаем друг друга. + +🟡 2 — Приятели +«Всегда подойду спросить: "Как дела? Как жизнь?"» +Видимся несколько раз в неделю. Общаемся и про учебу, и про жизнь, иногда затрагиваем личное (но не глубокое). Можем вместе пообедать или поиграть в пин-понг. + +🔵 3 — Знакомые +«Мы здороваемся в коридоре» +Видимся изредка, общение короткое и ситуативное. В основном на учебные/повседневные темы.""" + + +def make_type_str(type: str, profile: str, strong_sides: List[str], recomendation: str): + return f"""🎯ТИП: «{type}» + +📊 Ваш профиль: +{profile} + +💪 Ваши сильные стороны: +• {"\n• ".join(strong_sides)} + +🌟 Рекомендация: +{recomendation} +""" diff --git a/src/userdb.py b/src/userdb.py index 95e1d9b..8049cf0 100644 --- a/src/userdb.py +++ b/src/userdb.py @@ -1,12 +1,12 @@ import asyncio import os -from abc import ABC, abstractmethod from typing import List, Optional import pymongo.asynchronous.collection as pymongo_collection import pymongo.asynchronous.database as pymongo_database from dotenv import load_dotenv from pymongo import AsyncMongoClient +from pymongo.asynchronous.cursor import AsyncCursor from pymongo.errors import ServerSelectionTimeoutError from .models import Link, User, Username @@ -16,53 +16,26 @@ MONGODB_HOST = os.getenv("MONGODB_HOST") -class UserDB(ABC): - @abstractmethod - def add_user(self, user: User) -> None: - pass - - @abstractmethod - def get_user(self, username: Username) -> Optional[User]: - pass - - @abstractmethod - def get_links(self, username: Username) -> Optional[list[User]]: - pass - - @abstractmethod - def add_link(self, username: Username, link: Link) -> None: - pass - - class UserNotExist(Exception): pass -class MongoUserDB(UserDB): +class UserDB: client: Optional[AsyncMongoClient] = None db: Optional[pymongo_database.AsyncDatabase] = None collection: Optional[pymongo_collection.AsyncCollection] = None def __init__(self): - # Не создаем клиент в __init__, а делаем это лениво - pass - - async def _ensure_connection(self): - """Ленивая инициализация подключения к MongoDB""" - if self.client is None: - try: - self.client = AsyncMongoClient( - MONGODB_HOST, serverSelectionTimeoutMS=5000 - ) - self.db = self.client.get_database("cu_graph_bot") - self.collection = self.db.get_collection("users") - await self.collection.create_index("username", unique=True) - except ServerSelectionTimeoutError as e: - print("Ошибка подключения к MongoDB:", e) - raise e + try: + self.client = AsyncMongoClient(MONGODB_HOST, serverSelectionTimeoutMS=5000) + self.db = self.client.get_database("cu_graph_bot") + self.collection = self.db.get_collection("users") + asyncio.run(self.collection.create_index("username", unique=True)) + except ServerSelectionTimeoutError as e: + print("Ошибка подключения к MongoDB:", e) + raise e async def add_user(self, user: User) -> None: - await self._ensure_connection() current_user_to_add = await self.collection.find_one( {"username": user.username} ) @@ -72,7 +45,6 @@ async def add_user(self, user: User) -> None: pass async def get_links(self, username: Username) -> List[Link]: - await self._ensure_connection() all_friends = await self.collection.find_one({"username": username}) if all_friends is None: raise UserNotExist @@ -81,12 +53,16 @@ async def get_links(self, username: Username) -> List[Link]: return [Link.model_validate(i) for i in all_friends["_links"]] async def get_user(self, username: Username) -> Optional[User]: - await self._ensure_connection() all_user_data = await self.collection.find_one({"username": username}) return all_user_data + async def get_users(self, links_less_than: int) -> AsyncCursor: + all_user_data = self.collection.find( + {f"_links.{links_less_than}": {"$exists": False}} + ) + return all_user_data + async def add_link(self, username: Username, link: Link) -> None: - await self._ensure_connection() current_user = await self.collection.find_one({"username": username}) if current_user is None: raise UserNotExist @@ -102,42 +78,16 @@ async def add_link(self, username: Username, link: Link) -> None: {"username": username}, {"$push": {"_links": link.model_dump()}} ) - async def count_users(self) -> int: - await self._ensure_connection() + async def count_users(self, links_more_than: int = 4) -> int: count = len( - [i async for i in self.collection.find({"_links.4": {"$exists": True}})] + [ + i + async for i in self.collection.find( + {f"_links.{links_more_than}": {"$exists": True}} + ) + ] ) return count -class ListUserDB(UserDB): - _users: list[User] - - def __init__(self) -> None: - self._users = [] - - def add_user(self, user: User) -> None: - if user in self._users: - raise ValueError - self._users.append(user) - - def get_user(self, username: Username) -> Optional[User]: - for i in self._users: - if i.username == username: - return i - return None - - def get_links(self, username: Username) -> Optional[list[Link]]: - for i in self._users: - if i.username == username: - return i._links - return None - - def add_link(self, username: Username, link: Link) -> None: - for i in self._users: - if i.username == username: - i.set_link(link) - break - - -userdb = MongoUserDB() +userdb = UserDB() From 1bac2c95f7e4b29078aba37bfdbd6a40b8854f6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D0=BD=D0=B4=D1=80=D0=B5=D0=B9=20=D0=9C=D0=B8=D1=88?= =?UTF-8?q?=D1=83=D1=80=D0=B8=D0=BD?= Date: Thu, 4 Dec 2025 04:20:51 +0300 Subject: [PATCH 2/4] add: userid and chatid --- main.py | 6 +----- src/bot.py | 20 +++++++++++++++++++- src/models.py | 6 +++++- src/userdb.py | 10 ++++++++-- 4 files changed, 33 insertions(+), 9 deletions(-) diff --git a/main.py b/main.py index f0920d9..5194639 100644 --- a/main.py +++ b/main.py @@ -1,10 +1,6 @@ import asyncio -from src.bot import bot, dp - - -async def main(): - await dp.start_polling(bot) +from src.bot import main if __name__ == "__main__": asyncio.run(main()) diff --git a/src/bot.py b/src/bot.py index 4f6d5cf..3051571 100644 --- a/src/bot.py +++ b/src/bot.py @@ -18,7 +18,7 @@ from . import callbacks, templates from .models import Link, User, check_tg_username -from .userdb import userdb +from .userdb import UserDB load_dotenv() @@ -30,6 +30,12 @@ dp = Dispatcher() +async def main(): + global userdb + userdb = UserDB() + await dp.start_polling(bot) + + rkb = ReplyKeyboardMarkup( keyboard=[ [KeyboardButton(text="Мои контакты")], @@ -213,6 +219,9 @@ def rating_to_text(rating: int) -> str: @dp.message(F.text == "Мои контакты") async def get_usS(message: types.Message): + await userdb.add_ids_to_user( + message.from_user.username, message.from_user.id, message.chat.id + ) links = await userdb.get_links(message.from_user.username) if len(links) == 0: await message.answer("Ты ещё не добавил связи!\nВведи юзернейм (@username)") @@ -225,6 +234,9 @@ async def get_usS(message: types.Message): @dp.message(F.text == "Узнать тип личности") async def get_summary(message: types.Message): + await userdb.add_ids_to_user( + message.from_user.username, message.from_user.id, message.chat.id + ) links = await userdb.get_links(message.from_user.username) ratings = [i.rating for i in links] p1 = ratings.count(1) / len(ratings) @@ -312,6 +324,9 @@ async def get_summary(message: types.Message): @dp.message(F.text == "Кол-во пользователей") async def get_count(message: types.Message): + await userdb.add_ids_to_user( + message.from_user.username, message.from_user.id, message.chat.id + ) count = await userdb.count_users() await message.answer( f"Ботом уже воспользовались {count} человек{'а' if count % 10 >= 2 and count % 10 < 5 else ''}!\nНапоминаю, что для участия в розыгрыше нужно подписаться на @campusdna" @@ -320,6 +335,9 @@ async def get_count(message: types.Message): @dp.message(F.text[0] == "@") async def user_name_checker(message: types.Message): + userdb.add_ids_to_user( + message.from_user.username, message.from_user.id, message.chat.id + ) msg = (message.text).strip() try: username_to = check_tg_username(msg) diff --git a/src/models.py b/src/models.py index 985e847..8dda885 100644 --- a/src/models.py +++ b/src/models.py @@ -23,11 +23,15 @@ class Link(BaseModel): class User(BaseModel): + userid: int + chatid: int + username: Username - # TODO: Add sex, course, etc. + sex: Sex course: int = Field(ge=1, le=2) living: Living + _links: list[Link] = [] def set_link(self, link: Link): diff --git a/src/userdb.py b/src/userdb.py index 8049cf0..262d367 100644 --- a/src/userdb.py +++ b/src/userdb.py @@ -30,7 +30,7 @@ def __init__(self): self.client = AsyncMongoClient(MONGODB_HOST, serverSelectionTimeoutMS=5000) self.db = self.client.get_database("cu_graph_bot") self.collection = self.db.get_collection("users") - asyncio.run(self.collection.create_index("username", unique=True)) + asyncio.create_task(self.collection.create_index("username", unique=True)) except ServerSelectionTimeoutError as e: print("Ошибка подключения к MongoDB:", e) raise e @@ -89,5 +89,11 @@ async def count_users(self, links_more_than: int = 4) -> int: ) return count + # NOTE: This is temporary function for fixing current database + # And should be deleted someday because of obvious performance loss + async def add_ids_to_user(self, username: str, userid: int, chatid: int) -> None: + await self.collection.update_one( + {"username": username}, {"$set": {"userid": userid, "chatid": chatid}} + ) + -userdb = UserDB() From 6cf1cf95f705283bccf47083b43340b44d9ff788 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D0=BD=D0=B4=D1=80=D0=B5=D0=B9=20=D0=9C=D0=B8=D1=88?= =?UTF-8?q?=D1=83=D1=80=D0=B8=D0=BD?= Date: Thu, 4 Dec 2025 04:25:46 +0300 Subject: [PATCH 3/4] fix: next button --- src/bot.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/bot.py b/src/bot.py index 3051571..d93df3b 100644 --- a/src/bot.py +++ b/src/bot.py @@ -195,13 +195,21 @@ async def process_living( async def explaining_links(message: types.Message): await message.answer( templates.explaining_links_message, - callback_data=callbacks.TypeInfoCallback().pack(), + reply_markup=InlineKeyboardMarkup( + inline_keyboard=[ + [ + InlineKeyboardButton( + text="Далее", callback_data=callbacks.TypeInfoCallback().pack() + ) + ] + ] + ), ) -@dp.callback_query(callbacks.TypeInfoCallback) +@dp.callback_query(callbacks.TypeInfoCallback.filter()) async def start_survey(query: CallbackQuery, callback_data: callbacks.TypeInfoCallback): - await query.answer( + await query.message.answer( "Напиши юзернейм (@username) и я предложу тебе выбрать его категорию", reply_markup=rkb, ) @@ -381,7 +389,7 @@ async def user_name_checker(message: types.Message): async def process_data(query: CallbackQuery, callback_data: callbacks.LinkCallback): from_username = query.from_user.username if await userdb.get_user(from_username) is None: - await query.answer("Похоже тебе нужно перезапустить бота: /start") + await query.message.answer("Похоже тебе нужно перезапустить бота: /start") return await userdb.add_link( from_username, From a6d9b87994e420a73b807b6bbcc4cdb8c151c7b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D0=BD=D0=B4=D1=80=D0=B5=D0=B9=20=D0=9C=D0=B8=D1=88?= =?UTF-8?q?=D1=83=D1=80=D0=B8=D0=BD?= Date: Thu, 4 Dec 2025 04:46:44 +0300 Subject: [PATCH 4/4] refactor: make fsm storage more reliable --- src/bot.py | 12 +++++++++--- src/userdb.py | 24 +++++++----------------- 2 files changed, 16 insertions(+), 20 deletions(-) diff --git a/src/bot.py b/src/bot.py index d93df3b..621d890 100644 --- a/src/bot.py +++ b/src/bot.py @@ -6,6 +6,7 @@ from aiogram.filters import CommandStart from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup +from aiogram.fsm.storage.pymongo import PyMongoStorage from aiogram.types import ( CallbackQuery, InlineKeyboardButton, @@ -15,6 +16,7 @@ ) from aiogram.utils.formatting import as_list from dotenv import load_dotenv +from pymongo import AsyncMongoClient from . import callbacks, templates from .models import Link, User, check_tg_username @@ -26,13 +28,16 @@ if TOKEN is None: raise Exception("Couldn't find TG_BOT_TOKEN") +MONGODB_HOST = os.getenv("MONGODB_HOST") + +client = AsyncMongoClient(MONGODB_HOST) bot = Bot(token=TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML)) -dp = Dispatcher() +dp = Dispatcher(storage=PyMongoStorage(client, db_name="cu_graph_bot")) async def main(): - global userdb - userdb = UserDB() + global userdb, bot, dp + userdb = UserDB(client) await dp.start_polling(bot) @@ -189,6 +194,7 @@ async def process_living( living=callback_data.living, ) ) + await state.clear() await explaining_links(query.message) diff --git a/src/userdb.py b/src/userdb.py index 262d367..7fa169c 100644 --- a/src/userdb.py +++ b/src/userdb.py @@ -11,29 +11,19 @@ from .models import Link, User, Username -load_dotenv() - -MONGODB_HOST = os.getenv("MONGODB_HOST") - class UserNotExist(Exception): pass class UserDB: - client: Optional[AsyncMongoClient] = None - db: Optional[pymongo_database.AsyncDatabase] = None - collection: Optional[pymongo_collection.AsyncCollection] = None - - def __init__(self): - try: - self.client = AsyncMongoClient(MONGODB_HOST, serverSelectionTimeoutMS=5000) - self.db = self.client.get_database("cu_graph_bot") - self.collection = self.db.get_collection("users") - asyncio.create_task(self.collection.create_index("username", unique=True)) - except ServerSelectionTimeoutError as e: - print("Ошибка подключения к MongoDB:", e) - raise e + db: pymongo_database.AsyncDatabase + collection: pymongo_collection.AsyncCollection = None + + def __init__(self, client): + self.db = client.get_database("cu_graph_bot") + self.collection = self.db.get_collection("users") + asyncio.create_task(self.collection.create_index("username", unique=True)) async def add_user(self, user: User) -> None: current_user_to_add = await self.collection.find_one(