Add forced shift assignment functionality for workers and update tests

rota/shifts.py

@@ -1427,6 +1427,17 @@ class RotaBuilder(object):
             )
             logging.debug(f"Generate worker constraints: {worker.name}")
 
+            for week, day, shift_name in getattr(worker, "forced_assignments", []):
+                # Only add if this shift is valid for this week/day
+                if shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True):
+                    self.model.constraints.add(
+                        self.model.works[worker.id, week, day, shift_name] == 1
+                    )
+                else:
+                    self.add_warning(
+                        "Invalid forced assignment",
+                        f"Worker '{worker.name}' has forced assignment for non-existent shift: {shift_name} on {week}, {day}",
+                    )
 
             if hasattr(worker, "max_shifts_per_week_by_shift_name"):
                 for shift_type, max_per_week in worker.max_shifts_per_week_by_shift_name.items():
@@ -3182,13 +3193,18 @@ class RotaBuilder(object):
     def clear_shifts(self) -> None:
         self.shifts = []
 
-    def get_shift_names_by_week_day(self, week, day: DayStr) -> set:
+    def get_shift_names_by_week_day(self, week, day: DayStr, return_empty_if_week_day_not_found:bool=False) -> set:
         """Returns the shifts required for a specific day
 
         Returns:
             set: Set of ShiftName
 
         """
+        if (week, day) not in self.week_day_shifts_dict:
+            if return_empty_if_week_day_not_found:
+                return set()
+            raise WeekDayNotFound(f"Week {week}, Day {day} not found in week_day_shifts_dict")
+        #print(self.week_day_shifts_dict)
         return self.week_day_shifts_dict[(week, day)]
 
     def get_week_day_shift_names_by_week_day(self, week, day: DayStr) -> set:
@@ -3772,6 +3788,25 @@ class RotaBuilder(object):
         week = week - 1
         return self.start_date + datetime.timedelta(weeks=week)
 
+    def force_assign_shift(self, worker_name: str, week: int, day: str, shift_name: str):
+        """
+        Force a specific worker to be assigned to a shift on a specific week and day.
+        """
+        worker = self.get_worker_by_name(worker_name)
+        worker.forced_assignments.append((week, day, shift_name))
+
+    def force_assign_shift_by_date(self, worker_name: str, date: datetime.date, shift_name: str):
+        """
+        Force a specific worker to be assigned to a shift on a specific date.
+        """
+        # Find (week, day) for this date
+        for (week, day), dt in self.week_day_date_map.items():
+            if dt == date:
+                self.force_assign_shift(worker_name, week, day, shift_name)
+                return
+        raise ValueError(f"Date {date} is not in the rota period.")
+
+
     # RESULTS
     def export_rota_to_html(
         self, filename: str = "rota", folder=None, timestamp_filename=False
@@ -4494,3 +4529,8 @@ class WarningTermination(Exception):
     """Raised when a warning in the termination group is raised"""
 
     pass
+
+class WeekDayNotFound(Exception):
+    """Raised when a week/day combination is not found"""
+
+    pass

rota/workers.py

@@ -118,6 +118,7 @@ class Worker(BaseModel):
     hard_day_exclusions: list[tuple[str, str]] = []  # e.g. [("Fri", "Sun")]
     force_assign_with: dict[str, list[str]] = {}  # e.g. {"oncall": ["weekend"]}
     max_shifts_per_week_by_shift_name: dict[str, int] = {}  # e.g. {"night": 2, "a": 3}
+    forced_assignments: List[tuple[int, str, str]] = []  # (week, day, shift_name)
 
 
     assign_as_block_preferences: dict[str, float] = {}  # shift_name -> weight (positive = prefer block)
@@ -430,4 +431,8 @@ class Worker(BaseModel):
         """Helper to set the maximum number of shifts of a given type per week for this worker."""
         if not hasattr(self, "max_shifts_per_week_by_shift_name"):
             self.max_shifts_per_week_by_shift_name = {}
-        self.max_shifts_per_week_by_shift_name[shift_name] = max_per_week
+        self.max_shifts_per_week_by_shift_name[shift_name] = max_per_week
+
+    def force_assign_shift(self, week: int, day: str, shift_name: str):
+        """Force this worker to be assigned to a shift on a specific week/day."""
+        self.forced_assignments.append((week, day, shift_name))