From d7bcfce78c2c9a33362857b45c2bd538aacb79ea Mon Sep 17 00:00:00 2001 From: Ross Date: Sun, 10 Aug 2025 22:00:31 +0100 Subject: [PATCH] Implement hard day exclusions for workers and add corresponding tests --- gen_cons.py | 21 +++++++++---- rota/shifts.py | 72 +++++++++++++++++++++++++++++++++++++++++--- rota/workers.py | 5 +-- test/test_workers.py | 63 +++++++++++++++++++++++++++++++++++++- 4 files changed, 145 insertions(+), 16 deletions(-) diff --git a/gen_cons.py b/gen_cons.py index dd21537..7ebd7fa 100644 --- a/gen_cons.py +++ b/gen_cons.py @@ -266,7 +266,7 @@ def load_workers(): week=assignment["week"], day=assignment["day"], shift_name=assignment["shift"], - ) + ) if initial in rota_b_assignments: for date in rota_b_assignments.get(initial, []): w.force_assign_shift_by_date( @@ -285,6 +285,14 @@ def load_workers(): w.add_hard_day_dependency({"Fri", "Sat", "Sun"}, start_date="2025-11-17") else: w.add_hard_day_dependency({"Fri", "Sun"}, start_date="2025-11-17") + w.add_hard_day_exclusion({"Sat", "Sun"}, start_date="2025-11-17") + + + #if w.name == "TB": + # w.add_hard_day_exclusion({"Sat", "Sun"}, start_date="2025-11-17") + + + #w.set_max_shifts #w.set_max_shifts_per_week("weekend", 1) workers.append(w) @@ -316,6 +324,7 @@ def main( Rota.constraint_options["balance_weekends"] = True #Rota.constraint_options["max_weekend_frequency"] = 2 Rota.constraint_options["max_shifts_per_week"] = 6 + Rota.constraint_options["max_days_worked_per_week"] = 3 # Rota.constraint_options["avoid_st2_first_month"] = True Rota.add_shifts( @@ -342,7 +351,7 @@ def main( balance_offset=1, constraint=[ #{"name": "pre", "options": 2}, - #{"name": "post", "options": 1}, + {"name": "post", "options": 1}, ], display_char="a", #force_assign_with=["oncall"] @@ -371,10 +380,10 @@ def main( length=12.5, days=days[:5], balance_offset=1, - #constraint=[ - # {"name": "pre", "options": 1}, - # {"name": "post", "options": 1}, - #], + constraint=[ + #{"name": "pre", "options": 1}, + #{"name": "post", "options": 1}, + ], ), ) #Rota.allow_shifts_together_for_all_workers( diff --git a/rota/shifts.py b/rota/shifts.py index 4531d25..5ead3d6 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -565,6 +565,16 @@ class RotaBuilder(object): initialize=0, ) + self.model.works_day = Var( + ( + (worker.id, week, day) + for worker in self.workers + for week, day in self.get_week_day_combinations() + ), + within=Binary, + initialize=0, + ) + self.model.shift_week_worker_assigned = Var( ( (shift, week, worker.id) @@ -1422,7 +1432,7 @@ class RotaBuilder(object): ) # Most of our constraints apply per worker - # Worker constraint loop + # Worker constraint loop (worker loop) worker: Worker for worker in self.workers: console.rule( @@ -1430,6 +1440,47 @@ class RotaBuilder(object): ) logging.debug(f"Generate worker constraints: {worker.name}") + for week, day in self.get_week_day_combinations(): + shifts_today = self.get_shift_names_by_week_day(week, day) + total_assigned = sum( + self.model.works[worker.id, week, day, shift] + for shift in shifts_today + ) + # If assigned to any shift, works_day must be 1 + self.model.constraints.add(total_assigned >= self.model.works_day[worker.id, week, day]) + # If not assigned to any shift, works_day must be 0 + self.model.constraints.add(total_assigned <= 1000 * self.model.works_day[worker.id, week, day]) + + # --- Max days worked per week constraint (global or per-worker) --- + # Use per-worker value if set, otherwise use global constraint option if set + max_days_worked_per_week = getattr(worker, "max_days_worked_per_week", None) + if max_days_worked_per_week is None: + max_days_worked_per_week = self.constraint_options.get("max_days_worked_per_week", None) + if max_days_worked_per_week is not None: + for week in self.weeks: + # Create a binary variable for each day to indicate if the worker works any shift that day + for day in self.days: + var_name = f"works_any_{worker.id}_{week}_{day}" + if not hasattr(self.model, "works_any_day_vars"): + self.model.works_any_day_vars = {} + if (worker.id, week, day) not in self.model.works_any_day_vars: + self.model.works_any_day_vars[(worker.id, week, day)] = Var(within=Binary) + setattr(self.model, var_name, self.model.works_any_day_vars[(worker.id, week, day)]) + works_any = self.model.works_any_day_vars[(worker.id, week, day)] + total_assigned = sum( + self.model.works[worker.id, week, day, shift] + for shift in self.get_shift_names_by_week_day(week, day) + ) + # If assigned to any shift, works_any must be 1 + self.model.constraints.add(total_assigned >= works_any) + # If not assigned to any shift, works_any must be 0 + self.model.constraints.add(total_assigned <= 1000 * works_any) + # Now sum the binaries for the week + self.model.constraints.add( + sum(self.model.works_any_day_vars[(worker.id, week, day)] for day in self.days) + <= max_days_worked_per_week + ) + 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): @@ -1521,15 +1572,26 @@ class RotaBuilder(object): if hasattr(worker, "hard_day_exclusions"): for exclusion in worker.hard_day_exclusions: days_to_exclude = exclusion.days - weeks_to_apply = exclusion.weeks if exclusion.weeks is not None else self.weeks + weeks_to_apply = self.weeks + if exclusion.weeks is not None: + weeks_to_apply = exclusion.weeks + if exclusion.start_date is not None or exclusion.end_date is not None: + if exclusion.start_date is None: + exclusion.start_date = self.start_date + if exclusion.end_date is None: + exclusion.end_date = self.rota_end_date + + weeks_to_apply = self.get_weeks_by_date_range( + exclusion.start_date, exclusion.end_date + ) for week in weeks_to_apply: # For each week, prevent any shift being assigned on both days self.model.constraints.add( sum( - self.model.works[worker.id, week, day, shift] + self.model.works_day[worker.id, week, day] for day in days_to_exclude - for shift in self.get_shift_names_by_week_day(week, day) + #for shift in self.get_shift_names_by_week_day(week, day) ) <= 1 ) @@ -4222,7 +4284,7 @@ class RotaBuilder(object): if shift_name: shift = self.get_shift_by_name(shift_name) if "require_remote_site_presence_week" in shift.constraints: - remote_site = f" data-shift-remote-site='{shift.constraint_options['require_remote_site_presence_week'][0]}'" + remote_site = f" data-shift-remote-site='{shift.constraints['require_remote_site_presence_week'].options[0]}'" if "night" in shift.constraints: css_class = " ".join((css_class, "night-shift")) diff --git a/rota/workers.py b/rota/workers.py index 83192d7..ef33d5f 100644 --- a/rota/workers.py +++ b/rota/workers.py @@ -145,6 +145,7 @@ class Worker(BaseModel): 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) forced_assignments_by_date: List[tuple[datetime.date, str]] = [] # (date, shift_name) + max_days_worked_per_week: int | None= None # Default: no limit, can be set to a specific number assign_as_block_preferences: dict[str, float] = {} # shift_name -> weight (positive = prefer block) @@ -442,8 +443,6 @@ class Worker(BaseModel): days: List of days that must be grouped together. weeks: Optional list of weeks this dependency applies to. If None, applies to all weeks. """ - if not hasattr(self, "hard_day_dependencies"): - self.hard_day_dependencies = [] self.hard_day_dependencies.append(HardDayDependency(days=days, weeks=weeks, start_date=start_date, end_date=end_date)) def add_hard_day_exclusion(self, days: set[str], weeks: Optional[List[int]] = None, start_date: Optional[datetime.date] = None, end_date: Optional[datetime.date] = None): @@ -452,8 +451,6 @@ class Worker(BaseModel): days: List of days that must be excluded together. weeks: Optional list of weeks this exclusion applies to. If None, applies to all weeks. """ - if not hasattr(self, "hard_day_exclusions"): - self.hard_day_exclusions = [] self.hard_day_exclusions.append(HardDayExclusion(days=days, weeks=weeks, start_date=start_date, end_date=end_date)) def add_force_assign_with(self, shift: str, with_shifts: list[str] | str): diff --git a/test/test_workers.py b/test/test_workers.py index 0d7c048..4d3f471 100644 --- a/test/test_workers.py +++ b/test/test_workers.py @@ -1534,4 +1534,65 @@ def test_force_assign_shift2(): 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 + assert shift_list[3] == "a", "Worker 0 should have 'b' on 1st Tuesday" + + +def test_worker_hard_day_exclusion(): + Rota = generate_basic_rota(workers=4, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_hard_day_exclusion({"Mon", "Tue"}) + + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + 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_hard_day_exclusion", folder="tests") + + for worker in Rota.get_workers(): + shift_list = Rota.get_worker_shift_list(worker) + for week in range(16): + mon_idx = week * 7 + days.index("Mon") + tue_idx = week * 7 + days.index("Tue") + assert not (shift_list[mon_idx] == "a" and shift_list[tue_idx] == "a"), f"Worker {worker.name} has 'a' on both Mon and Tue in week {week}, which should not happen due to hard day exclusion" + +def test_worker_hard_day_exclusion2(): + Rota = generate_basic_rota(workers=4, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_hard_day_exclusion({"Sat", "Sun"}) + + worker.add_hard_day_dependency({"Fri", "Sun"}) + + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + length=8, + days=days[4:], + 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_hard_day_exclusion", folder="tests") + + for worker in Rota.get_workers(): + shift_list = Rota.get_worker_shift_list(worker) + for week in range(16): + sat_idx = week * 7 + days.index("Sat") + sun_idx = week * 7 + days.index("Sun") + assert not (shift_list[sat_idx] == "a" and shift_list[sun_idx] == "a"), f"Worker {worker.name} has 'a' on both Mon and Tue in week {week}, which should not happen due to hard day exclusion" \ No newline at end of file