diff --git a/daemons/fetch.py b/daemons/fetch.py index c1e3a16..ecdfa7a 100644 --- a/daemons/fetch.py +++ b/daemons/fetch.py @@ -1,16 +1,17 @@ import datetime from time import sleep -from helpers import now +from helpers import now, User +from helpers.models import UserSchema from helpers.mongo import mongo from helpers.ruz import ruz from settings import MOSCOW_TIMEZONE -def fetch_schedule_for_user(user_hse_id: int): +def fetch_schedule_for_user(user: User): today = now() next_day = today + datetime.timedelta(days=7) - schedule = ruz.get_schedule(user_hse_id, today, next_day) + schedule = ruz.get_schedule(user, today, next_day) if schedule is None: return False saved_ids = [] @@ -21,7 +22,7 @@ def fetch_schedule_for_user(user_hse_id: int): lesson = mongo.lessons_collection.find_one({ "discipline": element['discipline'], "auditorium": element['auditorium'], - "hse_user_id": user_hse_id, + "hse_user_id": user.hse_id, "begin": datetime.datetime( year=int(year), month=int(month), @@ -34,7 +35,7 @@ def fetch_schedule_for_user(user_hse_id: int): result = mongo.lessons_collection.insert_one({ "discipline": element['discipline'], "auditorium": element['auditorium'], - "hse_user_id": user_hse_id, + "hse_user_id": user.hse_id, "begin": datetime.datetime( year=int(year), month=int(month), @@ -56,13 +57,13 @@ def fetch_schedule_for_user(user_hse_id: int): saved_ids.append(result.inserted_id) else: saved_ids.append(lesson['_id']) - mongo.lessons_collection.delete_many({"hse_user_id": user_hse_id, "_id": {"$nin": saved_ids}}) + mongo.lessons_collection.delete_many({"hse_user_id": user.hse_id, "_id": {"$nin": saved_ids}}) return True def process(): for user in mongo.users_collection.find({"hse_id": {"$ne": None}}): - fetch_schedule_for_user(user['hse_id']) + fetch_schedule_for_user(UserSchema().load(user)) def delete_old(): diff --git a/daemons/notify.py b/daemons/notify.py index a34b01a..927491b 100644 --- a/daemons/notify.py +++ b/daemons/notify.py @@ -20,6 +20,7 @@ def process(): "notified": False }): ans = "" + ans += f"Дисциплина: {lesson['discipline']}\n" ans += f"Аудитория: {lesson['building']}, {lesson['auditorium']}\n" ans += f"Начало: {lesson['begin'].strftime('%H:%M')}\n" ans += f"Конец: {lesson['end'].strftime('%H:%M')}\n" diff --git a/helpers/answer.py b/helpers/answer.py index f74e88e..c23c48d 100644 --- a/helpers/answer.py +++ b/helpers/answer.py @@ -7,7 +7,7 @@ from daemons.bot import bot from daemons.fetch import fetch_schedule_for_user from helpers import now, get_next_daily_notify_time from helpers.keyboards import main_keyboard, notify_keyboard, yes_no_keyboard, again_keyboard, groups_keyboard, \ - no_daily_notify + no_daily_notify, student_or_teacher_keyboard from helpers.models import UserSchema, User from helpers.mongo import mongo from helpers.ruz import ruz @@ -36,9 +36,24 @@ class Answer(BaseAnswer): def handle_state_new(self, message: Message, user: User): bot.send_message( message.chat.id, - "Привет! Я буду помогать тебе выживать в вышке!\nДля начала пришли мне свое ФИО.", + "Привет! Я буду помогать тебе выживать в вышке!\nДля начала сообщи мне, ты преподаватель или студент?", + reply_markup=student_or_teacher_keyboard() ) - self.set_state(user, "wait_for_name") + self.set_state(user, "wait_for_student_or_teacher") + + def handle_state_wait_for_student_or_teacher(self, message: Message, user: User): + if message.text == "Студент": + bot.send_message(user.chat_id, "Принято! Теперь отправь мне свое ФИО.") + self.set_state(user, "wait_for_name") + elif message.text == "Преподаватель": + bot.send_message(user.chat_id, "Принято! Теперь отправь мне свое ФИО.") + mongo.users_collection.update_one( + {"chat_id": user.chat_id}, + {"$set": {"is_teacher": True, "state": "wait_for_name"}} + ) + self.set_state(user, "wait_for_name") + else: + bot.send_message(user.chat_id, "Ты отправил мне что-то неправильное, используй кнопки. Ты преподаватель или студент?", reply_markup=student_or_teacher_keyboard()) def handle_state_wait_for_name(self, message: Message, user: User): if message.text == "Начать заново": @@ -47,7 +62,8 @@ class Answer(BaseAnswer): "Привет! Я буду помогать тебе выживать в вышке!\nДля начала пришли мне свое ФИО.", ) return - data = ruz.find_person(message.text) + user.name = message.text + data = ruz.find_person(user) if data is None: bot.send_message( user.chat_id, @@ -55,22 +71,28 @@ class Answer(BaseAnswer): ) return if len(data) == 0: - bot.send_message(user.chat_id, "К сожалению, в РУЗе не нашлось такого студента, попробуй еще раз.") + bot.send_message(user.chat_id, "К сожалению, в РУЗе не нашлось такого человека, попробуй еще раз.") return - user.name = message.text mongo.users_collection.update_one( {"chat_id": user.chat_id}, {"$set": {"name": user.name}}) - bot.send_message( - user.chat_id, - "Отлично! Теперь выбери из списка свою группу.", - reply_markup=groups_keyboard(data) - ) + if user.is_teacher: + bot.send_message( + user.chat_id, + "Отлично! Теперь выбери из списка свой департамент.", + reply_markup=groups_keyboard(data) + ) + else: + bot.send_message( + user.chat_id, + "Отлично! Теперь выбери из списка свою группу.", + reply_markup=groups_keyboard(data) + ) self.set_state(user, "wait_for_group") def handle_state_wait_for_group(self, message: Message, user: User): group = message.text - data = ruz.find_person(user.name) + data = ruz.find_person(user) if data is None: bot.send_message( user.chat_id, @@ -84,8 +106,8 @@ class Answer(BaseAnswer): user.name = element['label'] break if user.group is None: - bot.send_message(user.chat_id, "Ты ввел что-то неправильно, попробуй еще раз сначала. Введи свое ФИО.") - self.set_state(user, "wait_for_name") + bot.send_message(user.chat_id, "Ты ввел что-то неправильно, попробуй еще раз сначала. Ты студент или преподаватель?", reply_markup=student_or_teacher_keyboard()) + self.set_state(user, "wait_for_student_or_teacher") return mongo.users_collection.update_one({"chat_id": user.chat_id}, {"$set": { "hse_id": user.hse_id, @@ -96,7 +118,7 @@ class Answer(BaseAnswer): user.chat_id, "Я нашел тебя в базе РУЗ. Я буду подсказывать тебе расписание, а также уведомлять о предстоящих парах.", ) - success = fetch_schedule_for_user(user.hse_id) + success = fetch_schedule_for_user(user) if success: lessons = mongo.get_today_lessons(user) if len(lessons) == 0: diff --git a/helpers/keyboards.py b/helpers/keyboards.py index b945eb2..69d022e 100644 --- a/helpers/keyboards.py +++ b/helpers/keyboards.py @@ -10,6 +10,13 @@ def main_keyboard(): return kb +def student_or_teacher_keyboard(): + kb = telebot.types.ReplyKeyboardMarkup(True, False) + kb.row("Преподаватель") + kb.row("Студент") + return kb + + def notify_keyboard(): kb = telebot.types.ReplyKeyboardMarkup(True, False) kb.row("Не уведомлять") diff --git a/helpers/models.py b/helpers/models.py index a5b083b..3e1927e 100644 --- a/helpers/models.py +++ b/helpers/models.py @@ -16,6 +16,7 @@ class User: notify_minutes: Optional[int] = 10 daily_notify_time: Optional[str] = None next_daily_notify_time: Optional[datetime.datetime] = None + is_teacher: bool = False class Meta: unknown = EXCLUDE diff --git a/helpers/ruz.py b/helpers/ruz.py index 635539c..3753ac6 100644 --- a/helpers/ruz.py +++ b/helpers/ruz.py @@ -3,7 +3,7 @@ import datetime from requests import get import settings - +from helpers import User fields = [ 'discipline', @@ -18,8 +18,12 @@ fields = [ class RUZ: - def find_person(self, name: str) -> dict | None: - search_str = settings.RUZ_API + f"search?term={name}&type=student" + def find_person(self, user: User) -> dict | None: + if user.is_teacher: + person_type = "person" + else: + person_type = "student" + search_str = settings.RUZ_API + f"search?term={user.name}&type={person_type}" try: data = get(search_str) except: @@ -28,10 +32,14 @@ class RUZ: return data.json() return None - def get_schedule(self, hse_id: int, begin_date: datetime.datetime, end_date: datetime.datetime): + def get_schedule(self, user: User, begin_date: datetime.datetime, end_date: datetime.datetime): start_date_str = begin_date.strftime("%Y.%m.%d") end_date_str = end_date.strftime("%Y.%m.%d") - search_str = settings.RUZ_API + f"schedule/student/{hse_id}?start={start_date_str}&finish={end_date_str}&lng=1" + if user.is_teacher: + person_type = "person" + else: + person_type = "student" + search_str = settings.RUZ_API + f"schedule/{person_type}/{user.hse_id}?start={start_date_str}&finish={end_date_str}&lng=1" try: data = get(search_str) except: