Implement hard day exclusions for workers and add corresponding tests

gen_cons.py

@@ -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=[
+            constraint=[
                 #{"name": "pre", "options": 1},
                 #{"name": "post", "options": 1},
-            #],
+            ],
         ),
     )
     #Rota.allow_shifts_together_for_all_workers(

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"))
 

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):

test/test_workers.py

@@ -1535,3 +1535,64 @@ def test_force_assign_shift2():
             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"
+
+
+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"