feat: add partial night shifts

Author: glichtner

execution_engine/omop/criterion/combination/temporal.py

@@ -15,6 +15,8 @@ class TimeIntervalType(StrEnum):
     MORNING_SHIFT = "morning_shift"
     AFTERNOON_SHIFT = "afternoon_shift"
     NIGHT_SHIFT = "night_shift"
+    NIGHT_SHIFT_BEFORE_MIDNIGHT = "night_shift_before_midnight"
+    NIGHT_SHIFT_AFTER_MIDNIGHT = "night_shift_after_midnight"
     DAY = "day"
     ANY_TIME = "any_time"
 
@@ -259,6 +261,34 @@ def NightShift(
 
         return cls.Presence(criterion, category, TimeIntervalType.NIGHT_SHIFT)
 
+    @classmethod
+    def NightShiftBeforeMidnight(
+        cls,
+        criterion: Union[Criterion, "CriterionCombination"],
+        category: CohortCategory,
+    ) -> "TemporalIndicatorCombination":
+        """
+        Create a NightShiftBeforeMidnight combination of criteria.
+        """
+
+        return cls.Presence(
+            criterion, category, TimeIntervalType.NIGHT_SHIFT_BEFORE_MIDNIGHT
+        )
+
+    @classmethod
+    def NightShiftAfterMidnight(
+        cls,
+        criterion: Union[Criterion, "CriterionCombination"],
+        category: CohortCategory,
+    ) -> "TemporalIndicatorCombination":
+        """
+        Create a NightShiftAfterMidnight combination of criteria.
+        """
+
+        return cls.Presence(
+            criterion, category, TimeIntervalType.NIGHT_SHIFT_AFTER_MIDNIGHT
+        )
+
     @classmethod
     def Day(
         cls,

execution_engine/settings.py

@@ -44,6 +44,12 @@ class TimeIntervalSettings(BaseModel):
     night_shift: TimeInterval = TimeInterval(
         start=datetime.time(22, 0, 0), end=datetime.time(5, 59, 59)
     )
+    night_shift_before_midnight: TimeInterval = TimeInterval(
+        start=datetime.time(22, 0, 0), end=datetime.time(23, 59, 59)
+    )
+    night_shift_after_midnight: TimeInterval = TimeInterval(
+        start=datetime.time(0, 0, 0), end=datetime.time(5, 59, 59)
+    )
     day: TimeInterval = TimeInterval(
         start=datetime.time(0, 0, 0), end=datetime.time(23, 59, 59)
     )

tests/execution_engine/omop/criterion/combination/test_temporal_combination.py

@@ -683,6 +683,88 @@ def patient_events(self, db_session, person_visit):
                 ),
                 {1: set(), 2: set(), 3: set()},
             ),
+            #######################
+            # Partial Night Shifts (before midnight)
+            #######################
+            (
+                TemporalIndicatorCombination.NightShiftBeforeMidnight(
+                    c1, category=CohortCategory.POPULATION
+                ),
+                {
+                    1: {
+                        (
+                            pendulum.parse("2023-03-01 22:00:00+01:00"),
+                            pendulum.parse("2023-03-01 23:59:59+01:00"),
+                        ),
+                        (
+                            pendulum.parse("2023-03-02 22:00:00+01:00"),
+                            pendulum.parse("2023-03-02 23:59:59+01:00"),
+                        ),
+                    },
+                    2: set(),
+                    3: set(),
+                },
+            ),
+            (
+                TemporalIndicatorCombination.NightShiftBeforeMidnight(
+                    c2, category=CohortCategory.POPULATION
+                ),
+                {
+                    1: {
+                        (
+                            pendulum.parse("2023-03-02 22:00:00+01:00"),
+                            pendulum.parse("2023-03-02 23:59:59+01:00"),
+                        ),
+                    },
+                    2: set(),
+                    3: set(),
+                },
+            ),
+            (
+                TemporalIndicatorCombination.NightShiftBeforeMidnight(
+                    c3, category=CohortCategory.POPULATION
+                ),
+                {1: set(), 2: set(), 3: set()},
+            ),
+            #######################
+            # Partial Night Shifts (after midnight)
+            #######################
+            (
+                TemporalIndicatorCombination.NightShiftAfterMidnight(
+                    c1, category=CohortCategory.POPULATION
+                ),
+                {
+                    1: {
+                        (
+                            pendulum.parse("2023-03-02 00:00:00+01:00"),
+                            pendulum.parse("2023-03-02 05:59:59+01:00"),
+                        ),
+                    },
+                    2: set(),
+                    3: set(),
+                },
+            ),
+            (
+                TemporalIndicatorCombination.NightShiftAfterMidnight(
+                    c2, category=CohortCategory.POPULATION
+                ),
+                {
+                    1: {
+                        (
+                            pendulum.parse("2023-03-03 00:00:00+01:00"),
+                            pendulum.parse("2023-03-03 05:59:59+01:00"),
+                        ),
+                    },
+                    2: set(),
+                    3: set(),
+                },
+            ),
+            (
+                TemporalIndicatorCombination.NightShiftAfterMidnight(
+                    c3, category=CohortCategory.POPULATION
+                ),
+                {1: set(), 2: set(), 3: set()},
+            ),
         ],
     )
     def test_combination_on_database(