Add forced shift assignment functionality for workers and update tests

.vscode/launch.json

@@ -20,5 +20,13 @@
             "program": "${workspaceFolder}/gen.py",
             "console": "integratedTerminal"
         },
+        {
+            "name": "Python: cons",
+            "type": "python",
+            "request": "launch",
+            "program": "${workspaceFolder}/gen_cons.py",
+            "console": "integratedTerminal"
+        },
+
     ]
 }

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)
@@ -431,3 +432,7 @@ class Worker(BaseModel):
         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
+
+    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))

test/test_workers.py

@@ -1456,3 +1456,82 @@ def test_max_shifts_per_week_by_shift_name_invalid_shift():
     )
     with pytest.raises(InvalidShift):
         Rota.build_and_solve(options={"ratio": 0.0})
+
+def test_force_assign_shift():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=2)
+
+    workers = Rota.get_workers()
+
+    workers[0].force_assign_shift(1, "Mon", "a")
+    workers[0].force_assign_shift(1, "Tue", "b")
+    workers[0].force_assign_shift(2, "Wed", "b")
+    workers[0].force_assign_shift(2, "Thu", "b")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign", folder="tests")
+
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker)
+        if worker == workers[0]:
+            # Worker 0 should have 'a' on Monday and 'b' on Tuesday
+            assert shift_list[0] == "a", "Worker 0 should have 'a' on 1st Monday"
+            assert shift_list[1] == "b", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[9] == "b", "Worker 0 should have 'b' on 2nd Wednesday"
+            assert shift_list[10] == "b", "Worker 0 should have 'b' on 2nd Thursday"
+
+def test_force_assign_shift2():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=1)
+
+    workers = Rota.get_workers()
+
+    workers[0].force_assign_shift(1, "Mon", "a")
+    workers[0].force_assign_shift(1, "Tue", "a")
+    workers[0].force_assign_shift(1, "Wed", "a")
+    workers[0].force_assign_shift(1, "Thu", "a")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign", folder="tests")
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker)
+        if worker == workers[0]:
+            # Worker 0 should have 'a' on Monday and 'b' on Tuesday
+            assert shift_list[0] == "a", "Worker 0 should have 'a' on 1st Monday"
+            assert shift_list[1] == "a", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[2] == "a", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[3] == "a", "Worker 0 should have 'b' on 1st Tuesday"