From 6bd47ba5b4cf856c8667bd410eb0f7d6c749f089 Mon Sep 17 00:00:00 2001 From: Ross Date: Fri, 8 Aug 2025 08:23:42 +0100 Subject: [PATCH] Add forced shift assignment functionality for workers and update tests --- .vscode/launch.json | 8 +++++ rota/shifts.py | 42 ++++++++++++++++++++++- rota/workers.py | 7 +++- test/test_workers.py | 81 +++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 135 insertions(+), 3 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index b604ae0..b37bd5e 100644 --- a/.vscode/launch.json +++ b/.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" + }, + ] } \ No newline at end of file diff --git a/rota/shifts.py b/rota/shifts.py index 493024a..1e498c3 100644 --- a/rota/shifts.py +++ b/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 \ No newline at end of file diff --git a/rota/workers.py b/rota/workers.py index 766cd5b..414e35b 100644 --- a/rota/workers.py +++ b/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 \ No newline at end of file + 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)) diff --git a/test/test_workers.py b/test/test_workers.py index 17540d1..953e56e 100644 --- a/test/test_workers.py +++ b/test/test_workers.py @@ -1455,4 +1455,83 @@ def test_max_shifts_per_week_by_shift_name_invalid_shift(): ), ) with pytest.raises(InvalidShift): - Rota.build_and_solve(options={"ratio": 0.0}) \ No newline at end of file + 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" \ No newline at end of file