Make next_lesson work with multiple validity ranges and year changes

aleksis/apps/chronos/managers.py

 from datetime import date, datetime, timedelta
 from enum import Enum
-from typing import Iterable, List, Optional, Union
+from typing import Dict, Iterable, List, Optional, Union
 
 from django.contrib.sites.managers import CurrentSiteManager as _CurrentSiteManager
 from django.db import models
@@ -10,7 +10,7 @@ from django.db.models.functions import Concat
 
 from calendarweek import CalendarWeek
 
-from aleksis.apps.chronos.util.date import week_weekday_from_date
+from aleksis.apps.chronos.util.date import week_weekday_from_date, week_weekday_to_date
 from aleksis.core.managers import DateRangeQuerySetMixin, SchoolTermRelatedQuerySet
 from aleksis.core.models import Group, Person
 from aleksis.core.util.core_helpers import get_site_preferences
@@ -375,28 +375,100 @@ class LessonDataQuerySet(models.QuerySet, WeekQuerySetMixin):
 
         return lesson_periods
 
-    def next_lesson(self, reference: "LessonPeriod", offset: Optional[int] = 1) -> "LessonPeriod":
+    def group_by_validity(self) -> Dict["ValidityRange", List["LessonPeriod"]]:
+        """Group lesson periods by validity range as dictionary."""
+        lesson_periods_by_validity = {}
+        for lesson_period in self:
+            lesson_periods_by_validity.setdefault(lesson_period.lesson.validity, [])
+            lesson_periods_by_validity[lesson_period.lesson.validity].append(lesson_period)
+        return lesson_periods_by_validity
+
+    def next_lesson(
+        self, reference: "LessonPeriod", offset: Optional[int] = 1
+    ) -> Optional["LessonPeriod"]:
         """Get another lesson in an ordered set of lessons.
 
         By default, it returns the next lesson in the set. By passing the offset argument,
         the n-th next lesson can be selected. By passing a negative number, the n-th
         previous lesson can be selected.
+
+        This function will handle week, year and validity range changes automatically
+        if the queryset contains enough lesson data.
         """
-        index = list(self.values_list("id", flat=True)).index(reference.id)
+        # Group lesson periods by validity to handle validity range changes correctly
+        lesson_periods_by_validity = self.group_by_validity()
+        validity_ranges = list(lesson_periods_by_validity.keys())
+
+        # List with lesson periods in the validity range of the reference lesson period
+        current_lesson_periods = lesson_periods_by_validity[reference.lesson.validity]
+        pks = [lesson_period.pk for lesson_period in current_lesson_periods]
+
+        # Position of the reference lesson period
+        index = pks.index(reference.id)
 
         next_index = index + offset
-        if next_index > self.count() - 1:
-            next_index %= self.count()
+        if next_index > len(pks) - 1:
+            next_index %= len(pks)
             week = reference._week + 1
         elif next_index < 0:
-            next_index = self.count() + next_index
+            next_index = len(pks) + next_index
             week = reference._week - 1
         else:
             week = reference._week
 
-        week = CalendarWeek(week=week, year=reference.lesson.get_year(week))
+        # Check if selected week makes a year change necessary
+        year = reference._year
+        if week < 1:
+            year -= 1
+            week = CalendarWeek.get_last_week_of_year(year)
+        elif week > CalendarWeek.get_last_week_of_year(year):
+            year += 1
+            week = 1
+
+        # Get the next lesson period in this validity range and it's date
+        # to check whether the validity range has to be changed
+        week = CalendarWeek(week=week, year=year)
+        next_lesson_period = current_lesson_periods[next_index]
+        next_lesson_period_date = week_weekday_to_date(week, next_lesson_period.period.period)
+
+        validity_index = validity_ranges.index(next_lesson_period.lesson.validity)
+
+        # If date of next lesson period is out of validity range (smaller) ...
+        if next_lesson_period_date < next_lesson_period.lesson.validity.date_start:
+            # ... we have to get the lesson period from the previous validity range
+            if validity_index == 0:
+                # There are no validity ranges (and thus no lessons)
+                # in the school term before this lesson period
+                return None
+
+            # Get new validity range and last lesson period of this validity range
+            new_validity = validity_ranges[validity_index - 1]
+            next_lesson_period = lesson_periods_by_validity[new_validity][-1]
+
+            # Build new week with the date from the new validity range/lesson period
+            week = CalendarWeek(
+                week=new_validity.date_end.isocalendar()[1], year=new_validity.date_end.year
+            )
+
+        # If date of next lesson period is out of validity range (larger) ...
+        elif next_lesson_period_date > next_lesson_period.lesson.validity.date_end:
+            # ... we have to get the lesson period from the next validity range
+            if validity_index >= len(validity_ranges):
+                # There are no validity ranges (and thus no lessons)
+                # in the school term after this lesson period
+                return None
+
+            # Get new validity range and first lesson period of this validity range
+            new_validity = validity_ranges[validity_index + 1]
+            next_lesson_period = lesson_periods_by_validity[new_validity][0]
+
+            # Build new week with the date from the new validity range/lesson period
+            week = CalendarWeek(
+                week=new_validity.date_start.isocalendar()[1], year=new_validity.date_start.year
+            )
 
-        return self.annotate_week(week).all()[next_index]
+        # Do a new query here to be able to annotate the new week
+        return self.annotate_week(week).get(pk=next_lesson_period.pk)
 
 
 class LessonPeriodQuerySet(LessonDataQuerySet, GroupByPeriodsMixin):