teacher

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():

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"

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:

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("Не уведомлять")

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

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: