To make data checks work with extra lessons/events, unify date querysets

aleksis/apps/alsijil/data_checks.py

 import logging
 
 from django.db.models import F
-from django.db.models.expressions import ExpressionWrapper, Func, Value
-from django.db.models.fields import DateField
-from django.db.models.functions import Concat
 from django.db.models.query_utils import Q
 from django.utils.translation import gettext as _
 
@@ -92,15 +89,6 @@ class NoGroupsOfPersonsSetInPersonalNotesDataCheck(DataCheck):
             cls.register_result(note)
 
 
-weekday_to_date = ExpressionWrapper(
-    Func(
-        Concat(F("year"), F("week")), Value("IYYYIW"), output_field=DateField(), function="TO_DATE"
-    )
-    + F("lesson_period__period__weekday"),
-    output_field=DateField(),
-)
-
-
 class LessonDocumentationOnHolidaysDataCheck(DataCheck):
     """Checks for lesson documentation objects on holidays.
 
@@ -123,13 +111,11 @@ class LessonDocumentationOnHolidaysDataCheck(DataCheck):
 
         holidays = Holiday.objects.all()
 
-        documentations = LessonDocumentation.objects.not_empty().annotate(
-            actual_date=weekday_to_date
-        )
+        documentations = LessonDocumentation.objects.not_empty().annotate_date_range()
 
         q = Q()
         for holiday in holidays:
-            q = q | Q(actual_date__gte=holiday.date_start, actual_date__lte=holiday.date_end)
+            q = q | Q(day_end__gte=holiday.date_start, day_start__lte=holiday.date_end)
         documentations = documentations.filter(q)
 
         for doc in documentations:
@@ -159,11 +145,11 @@ class PersonalNoteOnHolidaysDataCheck(DataCheck):
 
         holidays = Holiday.objects.all()
 
-        personal_notes = PersonalNote.objects.not_empty().annotate(actual_date=weekday_to_date)
+        personal_notes = PersonalNote.objects.not_empty().annotate_date_range()
 
         q = Q()
         for holiday in holidays:
-            q = q | Q(actual_date__gte=holiday.date_start, actual_date__lte=holiday.date_end)
+            q = q | Q(day_end__gte=holiday.date_start, day_start__lte=holiday.date_end)
         personal_notes = personal_notes.filter(q)
 
         for note in personal_notes:

aleksis/apps/alsijil/managers.py

 from datetime import date, datetime
 from typing import Optional, Sequence, Union
 
-from django.db.models import QuerySet
+from django.db.models import Case, ExpressionWrapper, F, Func, QuerySet, Value, When
+from django.db.models.fields import DateField
+from django.db.models.functions import Concat
 from django.db.models.query import Prefetch
 from django.db.models.query_utils import Q
 
@@ -11,6 +13,65 @@ from aleksis.apps.chronos.managers import DateRangeQuerySetMixin
 from aleksis.core.managers import CurrentSiteManagerWithoutMigrations
 
 
+class RegisterObjectRelatedQuerySet(QuerySet):
+    """Common queryset for personal notes and lesson documentations with shared API."""
+
+    def _get_weekday_to_date(self, weekday_name, year_name="year", week_name="week"):
+        """Get a ORM function which converts a weekday, a week and a year to a date."""
+        return ExpressionWrapper(
+            Func(
+                Concat(F(year_name), F(week_name)),
+                Value("IYYYIW"),
+                output_field=DateField(),
+                function="TO_DATE",
+            )
+            + F(weekday_name),
+            output_field=DateField(),
+        )
+
+    def annotate_day(self) -> QuerySet:
+        """Annotate every personal note/lesson documentation with the real date.
+
+        Attribute name: ``day``
+
+        .. note::
+            For events, this will annotate ``None``.
+        """
+        return self.annotate(
+            day=Case(
+                When(
+                    lesson_period__isnull=False,
+                    then=self._get_weekday_to_date("lesson_period__period__weekday"),
+                ),
+                When(
+                    extra_lesson__isnull=False,
+                    then=self._get_weekday_to_date(
+                        "extra_lesson__period__weekday", "extra_lesson__year", "extra_lesson__week"
+                    ),
+                ),
+            )
+        )
+
+    def annotate_date_range(self) -> QuerySet:
+        """Annotate every personal note/lesson documentation with the real date.
+
+        Attribute names: ``day_start``, ``day_end``
+
+        .. note::
+            For lesson periods and extra lessons,
+            this will annotate the same date for start and end day.
+        """
+        return self.annotate_day().annotate(
+            day_start=Case(
+                When(day__isnull=False, then="day"),
+                When(day__isnull=True, then="event__date_start"),
+            ),
+            day_end=Case(
+                When(day__isnull=False, then="day"), When(day__isnull=True, then="event__date_end"),
+            ),
+        )
+
+
 class PersonalNoteManager(CurrentSiteManagerWithoutMigrations):
     """Manager adding specific methods to personal notes."""
 
@@ -33,7 +94,7 @@ class PersonalNoteManager(CurrentSiteManagerWithoutMigrations):
         )
 
 
-class PersonalNoteQuerySet(QuerySet):
+class PersonalNoteQuerySet(RegisterObjectRelatedQuerySet, QuerySet):
     def not_empty(self):
         """Get all not empty personal notes."""
         return self.filter(
@@ -45,7 +106,7 @@ class LessonDocumentationManager(CurrentSiteManagerWithoutMigrations):
     pass
 
 
-class LessonDocumentationQuerySet(QuerySet):
+class LessonDocumentationQuerySet(RegisterObjectRelatedQuerySet, QuerySet):
     def not_empty(self):
         """Get all not empty lesson documentations."""
         return self.filter(~Q(topic="") | ~Q(group_note="") | ~Q(homework=""))