diff --git a/aleksis/apps/chronos/model_extensions.py b/aleksis/apps/chronos/model_extensions.py index 917a5dbad3e6e0ddfd60d6c052c15314ff937e4d..2e092d2589481a1712ac521abb2b8ca69282a936 100644 --- a/aleksis/apps/chronos/model_extensions.py +++ b/aleksis/apps/chronos/model_extensions.py @@ -81,35 +81,44 @@ def lesson_periods_as_teacher(self): @Person.method -def daily_lessons(self, day: date): - """Get all lessons of this person on the given day.""" - return LessonPeriod.objects.on_day(day).filter_from_person(self) +def lessons_on_day(self, day: date): + """Get all lessons of this person (either as participant or teacher) on the given day.""" + return ( + LessonPeriod.objects.on_day(day) + .filter_from_person(self) + .order_by("period__period") + ) @Person.method -def next_lesson(self, lesson_period: "LessonPeriod", day: date) -> Union["LessonPeriod", None]: - """Get next lesson of the person on the same day.""" - daily_lessons = self.daily_lessons(day) +def _adjacent_lesson( + self, lesson_period: "LessonPeriod", day: date, offset: int = 1 +) -> Union["LessonPeriod", None]: + """Get next/previous lesson of the person (either as participant or teacher) on the same day.""" + daily_lessons = self.lessons_on_day(day) ids = list(daily_lessons.values_list("id", flat=True)) index = ids.index(lesson_period.pk) - if index + 1 < len(ids): - return daily_lessons[index + 1] + if (offset > 0 and index + offset < len(ids)) or (offset < 0 and index >= -offset): + return daily_lessons[index + offset] else: return None @Person.method -def previous_lesson(self, lesson_period: "LessonPeriod", day: date) -> Union["LessonPeriod", None]: - """Get previous lesson of the person on the same day.""" - daily_lessons = self.daily_lessons(day) - ids = list(daily_lessons.values_list("id", flat=True)) - index = ids.index(lesson_period.pk) +def next_lesson( + self, lesson_period: "LessonPeriod", day: date +) -> Union["LessonPeriod", None]: + """Get next lesson of the person (either as participant or teacher) on the same day.""" + return self._adjacent_lesson(lesson_period, day) - if index > 0: - return daily_lessons[index - 1] - else: - return None + +@Person.method +def previous_lesson( + self, lesson_period: "LessonPeriod", day: date +) -> Union["LessonPeriod", None]: + """Get previous lesson of the person (either as participant or teacher) on the same day.""" + return self._adjacent_lesson(lesson_period, day, offset=-1) def for_timetables(cls):