90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
import datetime
|
|
from functools import cached_property
|
|
|
|
import pymongo
|
|
|
|
import settings
|
|
from helpers import now
|
|
from helpers.models import UserSchema, User
|
|
|
|
|
|
class Mongo:
|
|
def __init__(self):
|
|
url = f"mongodb://{settings.MONGO_USER}:{settings.MONGO_PASSWORD}@{settings.MONGO_HOST}:27017/"
|
|
self.client = pymongo.MongoClient(url)
|
|
self.database = self.client.get_database("ruz-bot")
|
|
self.users_collection.create_index([
|
|
("chat_id", 1)
|
|
])
|
|
self.users_collection.create_index([
|
|
("notify_minutes", 1)
|
|
])
|
|
self.users_collection.create_index([
|
|
("next_daily_notify_time", 1)
|
|
])
|
|
self.lessons_collection.create_index([
|
|
("discipline", 1),
|
|
("auditorium", 1),
|
|
("begin", 1),
|
|
("hse_user_id", 1),
|
|
("link", 1)
|
|
])
|
|
self.lessons_collection.create_index([
|
|
("hse_user_id", 1),
|
|
("begin", 1)
|
|
])
|
|
self.lessons_collection.create_index([
|
|
("hse_user_id", 1),
|
|
("begin", 1),
|
|
("notified", 1)
|
|
])
|
|
|
|
def get_user(self, user_id: int) -> User:
|
|
return UserSchema().loads(self.users_collection.find_one({"id": user_id}))
|
|
|
|
def __getitem__(self, item):
|
|
return self.database.get_collection(item)
|
|
|
|
@cached_property
|
|
def users_collection(self):
|
|
return self["users"]
|
|
|
|
@cached_property
|
|
def lessons_collection(self):
|
|
return self["lessons"]
|
|
|
|
def _get_lessons_on_date(self, user: User, date: datetime.datetime):
|
|
date = datetime.datetime(year=date.year, month=date.month, day=date.day)
|
|
next_date = date + datetime.timedelta(days=1)
|
|
lessons = []
|
|
for lesson in self.lessons_collection.find({
|
|
"hse_user_id": user.hse_id,
|
|
"begin": {"$gte": date, "$lte": next_date}}
|
|
):
|
|
lessons.append(lesson)
|
|
lessons.sort(key=lambda les: les["begin"])
|
|
return lessons
|
|
|
|
def get_today_lessons(self, user: User):
|
|
return self._get_lessons_on_date(user, now(user))
|
|
|
|
def get_tomorrow_lessons(self, user: User):
|
|
return self._get_lessons_on_date(user, now(user) + datetime.timedelta(days=1))
|
|
|
|
def get_week_lessons(self, user: User):
|
|
date = now(user)
|
|
date = datetime.datetime(year=date.year, month=date.month, day=date.day)
|
|
weekday = date.weekday()
|
|
next_date = date + datetime.timedelta(days=(6 - weekday))
|
|
lessons = []
|
|
for lesson in self.lessons_collection.find({
|
|
"hse_user_id": user.hse_id,
|
|
"begin": {"$gte": date, "$lte": next_date}}
|
|
):
|
|
lessons.append(lesson)
|
|
lessons.sort(key=lambda les: les["begin"])
|
|
return lessons
|
|
|
|
|
|
mongo = Mongo()
|