From fc3102ec4723fe7d864e63de010b97fa9b262d04 Mon Sep 17 00:00:00 2001 From: Ross Date: Sun, 22 Jun 2025 09:38:03 +0100 Subject: [PATCH] numberous fixes and improvements --- rota/shifts.py | 130 +++++++++++++++++++++----- rota/workers.py | 21 ++++- test/test_workers.py | 212 ++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 333 insertions(+), 30 deletions(-) diff --git a/rota/shifts.py b/rota/shifts.py index 8f1783e..808c478 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -145,6 +145,7 @@ class SingleShift(BaseModel): end_date: datetime.date | None = None pair_proxy: bool = False display_char: str | None = None # Add this line + force_assign_with: list[str] = [] # List of shift names to force assign with this shift on the same day model_config = ConfigDict( extra="allow", @@ -1391,17 +1392,39 @@ class RotaBuilder(object): ) logging.debug(f"Generate worker constraints: {worker.name}") + for week, day, shift in self.get_all_shiftclass_combinations(): + # Only apply if this shift has force_assign_with set + if hasattr(shift, "force_assign_with") and shift.force_assign_with: + for other_shift_name in shift.force_assign_with: + # Only add constraint if the other shift is available on this day + if other_shift_name in self.get_shift_names_by_week_day(week, day): + for worker in self.get_workers_for_shift(shift): + self.model.constraints.add( + self.model.works[worker.id, week, day, shift.name] + <= self.model.works[worker.id, week, day, other_shift_name] + ) + # Add hard shift dependencies for this worker if hasattr(worker, "hard_day_dependencies"): - for from_day, to_day in worker.hard_day_dependencies.items(): + for day_group in worker.hard_day_dependencies: for week in self.weeks: - for shift in self.get_shift_names_by_week_day(week, from_day): - # Only apply if the shift also exists on the to_day - if shift in self.get_shift_names_by_week_day(week, to_day): - self.model.constraints.add( - self.model.works[worker.id, week, from_day, shift] - <= self.model.works[worker.id, week, to_day, shift] - ) + assigned_any_shift = {} + for day in day_group: + shifts_today = self.get_shift_names_by_week_day(week, day) + var_name = f"hard_day_dep_{worker.id}_{week}_{day}" + if not hasattr(self.model, "hard_day_dep_vars"): + self.model.hard_day_dep_vars = {} + if (worker.id, week, day) not in self.model.hard_day_dep_vars: + self.model.hard_day_dep_vars[(worker.id, week, day)] = Var(within=Binary) + setattr(self.model, var_name, self.model.hard_day_dep_vars[(worker.id, week, day)]) + assigned_any_shift[day] = self.model.hard_day_dep_vars[(worker.id, week, day)] + total_assigned = sum(self.model.works[worker.id, week, day, shift] for shift in shifts_today) + self.model.constraints.add(total_assigned >= assigned_any_shift[day]) + self.model.constraints.add(total_assigned <= len(shifts_today) * assigned_any_shift[day]) + # All-or-none: all days in the group must have the same assignment status + vals = list(assigned_any_shift.values()) + for i in range(1, len(vals)): + self.model.constraints.add(vals[i] == vals[0]) # Add hard day exclusions for this worker if hasattr(worker, "hard_day_exclusions"): @@ -1448,6 +1471,17 @@ class RotaBuilder(object): - len(allowed_set) + 1 ) + shifts_today = self.get_shift_names_by_week_day(week, day) + for shift_name in shifts_today: + if hasattr(worker, "force_assign_with") and shift_name in worker.force_assign_with: + for other_shift in worker.force_assign_with[shift_name]: + if other_shift in shifts_today: + self.model.constraints.add( + self.model.works[worker.id, week, day, shift_name] + <= self.model.works[worker.id, week, day, other_shift] + ) + + if self.get_workers_who_require_locums(): for week, day, shift in self.get_all_shiftclass_combinations(): @@ -2215,17 +2249,20 @@ class RotaBuilder(object): # for day in self.days[5:] # for shift in self.get_shifts()) / 2) # if self.constraint_options["balance_weekends"]: - self.model.constraints.add( - self.model.works_weekend[worker.id, week] - >= sum( - self.model.works[worker.id, week, day, shiftname] - for w, day, shiftname in self.get_all_shiftname_combinations( - week=week - ) - if day in self.days[5:] - ) - / 2 - ) + + + # Try disabling tihs to allow more than 2 shifts ot be assigned on a weekend + #self.model.constraints.add( + # self.model.works_weekend[worker.id, week] + # >= sum( + # self.model.works[worker.id, week, day, shiftname] + # for w, day, shiftname in self.get_all_shiftname_combinations( + # week=week + # ) + # if day in self.days[5:] + # ) + # / 2 + #) self.model.constraints.add( self.model.works_weekend[worker.id, week] <= sum( @@ -2492,11 +2529,24 @@ class RotaBuilder(object): worker_allowed_sets = getattr( worker, "allowed_multi_shift_sets", [] ) + force_assign_sets = [] + for shift_name in shifts_today: + if hasattr(worker, "force_assign_with") and shift_name in worker.force_assign_with: + force_set = set([shift_name] + list(worker.force_assign_with[shift_name])) + if force_set.issubset(shifts_today): + force_assign_sets.append(force_set) + shift_obj = self.get_shift_by_name(shift_name) + if getattr(shift_obj, "force_assign_with", []): + force_set = set([shift_name] + list(shift_obj.force_assign_with)) + if force_set.issubset(shifts_today): + force_assign_sets.append(force_set) + + # If no allowed sets, we can only assign one shift allowed_sets_today = [ allowed_set for allowed_set in worker_allowed_sets if allowed_set.issubset(shifts_today) - ] + ] + force_assign_sets if allowed_sets_today: # Forbid any multi-shift assignment that includes shifts from different allowed sets @@ -3182,6 +3232,24 @@ class RotaBuilder(object): else: s.end_date = self.rota_end_date + + if hasattr(s, "force_assign_with") and s.force_assign_with: + missing = set(s.force_assign_with) - set(shift.name for shift in self.shifts) + if missing: + raise InvalidShift( + f"Shift '{s.name}' has force_assign_with referencing non-existent shift(s): {missing}" + ) + + all_shift_names = set(shift.name for shift in self.shifts) + for worker in self.workers: + if hasattr(worker, "force_assign_with"): + for shift_name, with_shifts in getattr(worker, "force_assign_with", {}).items(): + missing = set([shift_name] + list(with_shifts)) - all_shift_names + if missing: + raise InvalidShift( + f"Worker '{worker.name}' has force_assign_with referencing non-existent shift(s): {missing}" + ) + # Check worker requirements if not isinstance(s.workers_required, int): # Validate the worker requirements @@ -4166,22 +4234,36 @@ class RotaBuilder(object): return shifts - def get_worker_shift_list(self, worker: Worker, include_locums=False) -> List: + def get_worker_shift_list(self, worker: Worker, include_locums=False, search_multiple_assignments=False) -> List: + """ + This function returns a list of shifts assigned to a worker for each week and day. + + It was originally designed to return a list of shift names (when only a single shift per day was supported), but it can also return a set of shift names if `search_multiple_assignments` is True. + """ shifts = [] for week, day in self.get_week_day_combinations(): # d = self.start_date + datetime.timedelta(n) # n = n + 1 - temp = "" + if search_multiple_assignments: + temp = set() + else: + temp = "" for shift in self.get_shift_names_by_week_day(week, day): if self.model.works[worker.id, week, day, shift].value > 0.90: - temp = shift + if search_multiple_assignments: + temp.add(shift) + else: + temp = shift elif ( include_locums and self.model.locum_works[worker.id, week, day, shift].value > 0.90 ): - temp = shift + if search_multiple_assignments: + temp.add(shift) + else: + temp = shift shifts.append(temp) return shifts diff --git a/rota/workers.py b/rota/workers.py index 93f70c6..47b4bb7 100644 --- a/rota/workers.py +++ b/rota/workers.py @@ -114,13 +114,17 @@ class Worker(BaseModel): allowed_multi_shift_sets: list[set] = [] # or list/set of frozensets prefer_multi_shift_together: int = 0 # Default: no preference # Set to a positive integer to prefer, negative to discourage, 0 for neutral - hard_day_dependencies: dict[str, str] = {} # e.g. {"Fri": "Sun"} + hard_day_dependencies: set[frozenset[str]] = set() # e.g. {frozenset({"Mon", "Tue"}), frozenset({"Fri", "Sun"})} hard_day_exclusions: list[tuple[str, str]] = [] # e.g. [("Fri", "Sun")] + force_assign_with: dict[str, list[str]] = {} # e.g. {"oncall": ["weekend"]} + assign_as_block_preferences: dict[str, float] = {} # shift_name -> weight (positive = prefer block) shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist + weekend_shift_target_number: int = 0 + model_config = ConfigDict( extra="allow", ) @@ -402,12 +406,19 @@ class Worker(BaseModel): self.allowed_multi_shift_sets = [] self.allowed_multi_shift_sets.append(frozenset(shifts)) - def add_hard_day_dependency(self, from_day: str, to_day: str): + def add_hard_day_dependency(self, *days: str): if not hasattr(self, "hard_day_dependencies"): - self.hard_day_dependencies = {} - self.hard_day_dependencies[from_day] = to_day + self.hard_day_dependencies = set() + self.hard_day_dependencies.add(frozenset(days)) def add_hard_day_exclusion(self, day1: str, day2: str): if not hasattr(self, "hard_day_exclusions"): self.hard_day_exclusions = [] - self.hard_day_exclusions.append((day1, day2)) \ No newline at end of file + self.hard_day_exclusions.append((day1, day2)) + + def add_force_assign_with(self, shift: str, with_shifts: list[str] | str): + if not hasattr(self, "force_assign_with"): + self.force_assign_with = {} + if isinstance(with_shifts, str): + with_shifts = [with_shifts] + self.force_assign_with[shift] = with_shifts \ No newline at end of file diff --git a/test/test_workers.py b/test/test_workers.py index c16cc67..33b4895 100644 --- a/test/test_workers.py +++ b/test/test_workers.py @@ -1,5 +1,5 @@ import pytest -from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days +from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days import datetime from rota.workers import NonWorkingDays, NotAvailableToWork, Worker, generate_not_available_to_works @@ -1169,4 +1169,214 @@ def test_worker_hard_day_dependency_weekend2(): assert shift_list[sat_idx] == "a", f"Worker4 assigned 'a' on Fri (week {week}) but not on Sun" +def test_worker_force_assign_with_invalid_shift(): + Rota = generate_basic_rota(workers=4, weeks_to_rota=2) + workers = Rota.get_workers() + for worker in workers: + worker.add_force_assign_with("a", "c") + + 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 + ), + ) + + with pytest.raises(InvalidShift): + # This should raise an InvalidShift error because 'c' does not exist + Rota.build_and_solve(options={"ratio": 0.0}) + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="c", + length=8, + days=days[:4], + workers_required=1, + assign_as_block=False, # global setting is off + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + +def test_worker_force_assign_with(): + Rota = generate_basic_rota(workers=4, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_force_assign_with("a", "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_with", folder="tests") + + for worker in Rota.get_workers(): + shift_list = Rota.get_worker_shift_list(worker, search_multiple_assignments=True) + + for shift in shift_list: + if shift == "a": + assert "b" in shift_list, f"Worker {worker.name} has 'a' but not 'b' assigned with it" + elif shift == "b": + assert "a" in shift_list, f"Worker {worker.name} has 'b' but not 'a' assigned with it" + + assert shift_list.count({"a", "b"}) == 16 + +def test_worker_force_assign_with_and_hard_days(): + Rota = generate_basic_rota(workers=4, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_force_assign_with("a", "b") + + #worker.add_hard_day_dependency("Sat", "Sun") + worker.add_hard_day_dependency("Mon", "Tue") + + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + length=8, + days=days[:2], + workers_required=1, + assign_as_block=False, # global setting is off + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[:2], + 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_with_hard_days", folder="tests") + + assert Rota.results.solver.status == "ok" + +def test_worker_force_assign_with_and_hard_days_weekend(): + Rota = generate_basic_rota(workers=4, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_force_assign_with("a", "b") + + #worker.add_hard_day_dependency("Sat", "Sun") + worker.add_hard_day_dependency("Sat", "Sun") + + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + length=8, + days=days[5:], + workers_required=1, + assign_as_block=False, # global setting is off + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[5:], + 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_with_hard_days", folder="tests") + + assert Rota.results.solver.status == "ok" + + +def test_workers_double_shifts_on_mon_tue(): + Rota = generate_basic_rota(workers=2, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_force_assign_with("a", "b") + + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + length=8, + days=days[:2], + workers_required=2, + assign_as_block=False, # global setting is off + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[:2], + workers_required=2, + assign_as_block=False, # global setting is off + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + Rota.export_rota_to_html("test_worker_double_shifts", folder="tests") + + assert Rota.results.solver.status == "ok" + +def test_workers_double_shifts_on_weekends(): + Rota = generate_basic_rota(workers=2, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_force_assign_with("a", "b") + + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + length=8, + days=days[5:], + workers_required=2, + assign_as_block=False, # global setting is off + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[5:], + workers_required=2, + assign_as_block=False, # global setting is off + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + Rota.export_rota_to_html("test_worker_double_shifts", folder="tests") + + assert Rota.results.solver.status == "ok" \ No newline at end of file