Add max days per week block constraint and exclude days for shift constraints

gen_cons.py

@@ -3,7 +3,7 @@ import datetime
 import os
 import sys
 import time
-from rota.shifts import NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, WorkerRequirement, days
+from rota.shifts import MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, WorkerRequirement, days
 import typer
 import subprocess
 import pandas as pd
@@ -336,6 +336,7 @@ def main(
     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.constraint_options["max_days_per_week_block"] = [(4, 2), (6, 3)]  # (max_days, week_block)
 
     #Rota.enable_unavailabilities = False
 
@@ -348,6 +349,8 @@ def main(
             balance_offset=2,
             assign_as_block=False,
             constraints=[
+                PreShiftConstraint(days=2, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun")),
+                PostShiftConstraint(days=1, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun")),
                 #{
                 #    "name": "max_shifts_per_week",
                 #    "options": 3,
@@ -395,7 +398,10 @@ def main(
             days=days[:5],
             balance_offset=1,
             constraints=[
-                PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"]),
+                #PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"]),
+                PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"], exclude_days=("Sat", "Sun")),
+                #PostShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall", "weekend a", "weekend b"], include_weekends=False),
+                MaxShiftsPerWeekConstraint(max_shifts=1),
                 #{"name": "post", "options": 1},
             ],
         ),

rota/shifts.py

@@ -92,10 +92,12 @@ class NightConstraint(BaseShiftConstraint):
 class PreShiftConstraint(BaseShiftConstraint):
     days: Any = None
     ignore_shifts: List[str] = []
+    exclude_days: List[str] = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
 
 class PostShiftConstraint(BaseShiftConstraint):
     days: Any = None
     ignore_shifts: List[str]  = []
+    exclude_days: List[str] = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
 
 class BalanceAcrossGroupsConstraint(BaseShiftConstraint):
     """"""
@@ -385,6 +387,7 @@ class RotaBuilder(object):
             "max_night_frequency_week_exclusions": [],
             "max_shifts_per_week": 7,
             "max_shifts_per_month": 40,
+            "max_days_per_week_block": [(3, 2)],  # (max_days, week_block)
             # The following will cause an unsolvable problem if a shift
             # is forced as a block and spans the time peroid
             "prevent_monday_after_full_weekends": [],
@@ -1840,6 +1843,17 @@ class RotaBuilder(object):
                     >= -self.constraint_options["maximum_allowed_shift_diff"]
                 )
 
+            for max_days, weeks in self.constraint_options["max_days_per_week_block"]:
+                for week_blocks in self.get_week_block_iterator(weeks):
+                    self.model.constraints.add(
+                        sum(
+                            self.model.works_day[worker.id, week, day]
+                            for week, day in self.get_week_day_combinations()
+                            if week in week_blocks
+                        )
+                        <= max_days
+                    )
+
             for week_blocks in self.get_week_block_iterator(4):
                 # Prevent more than n number shifts per 4 weeks
                 try:
@@ -2922,6 +2936,8 @@ class RotaBuilder(object):
                     if week not in constraint.weeks:
                         continue
                     for n in range(0, constraint.days):
+                        if day in constraint.exclude_days:
+                            continue
                         if day in constraint_shift.days:
                             # Only apply if worker can work this shift
                             if worker.site not in constraint_shift.sites:
@@ -2957,6 +2973,8 @@ class RotaBuilder(object):
                     if week not in constraint.weeks:
                         continue
                     for n in range(0, constraint.days):
+                        if day in constraint.exclude_days:
+                            continue
                         if day in constraint_shift.days:
                             # Only apply if worker can work this shift
                             if worker.site not in constraint_shift.sites: