a few more fixes and improvements

gen_cons.py

@@ -29,7 +29,7 @@ def main(
     time_to_run: int = 60 * 60,
     ratio: float = 0.001,
     start_date: datetime.datetime = "2025-09-01",
-    weeks: int = 16,
+    weeks: int = 12,
     bom: int = 1,
 ):
     rota_start_date = start_date.date()
@@ -42,61 +42,42 @@ def main(
         name="cons rota",
     )
     # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
-    Rota.constraint_options["balance_shifts"] = False
-    Rota.constraint_options["balance_shifts_quadratic"] = True
     Rota.constraint_options["balance_weekends"] = True
-    Rota.constraint_options["max_night_frequency"] = 3  # 3
-    Rota.constraint_options["max_night_frequency_week_exclusions"] = []
     Rota.constraint_options["max_weekend_frequency"] = 4
-    Rota.constraint_options["hard_constrain_pair_separation"] = True
+    Rota.constraint_options["max_shifts_per_week"] = 20
-    Rota.constraint_options["max_shifts_per_week"] = 3
     # Rota.constraint_options["avoid_st2_first_month"] = True
 
     Rota.add_shifts(
         SingleShift(
             sites=("rota a",),
-            name="oncall weekday",
+            name="oncall",
             length=12.5,
-            days=days[:4],
+            days=days,
-            balance_offset=2,
+            balance_offset=5,
+            assign_as_block=False,
             constraint=[
-                {
+                #{
-                    "name": "max_shifts_per_week",
+                #    "name": "max_shifts_per_week",
-                    "options": 3,
+                #    "options": 3,
-                }
+                #}
             ],
         ),
         SingleShift(
             sites=("rota a",),
-            name="Oncall weekend",
+            name="weekend",
-            length=12.5,
-            days=days[4:],
-            balance_offset=2,
-            assign_as_block=True,
-            constraint=[
-                {
-                    "name": "max_shifts_per_week",
-                    "options": 3,
-                },
-                {"name": "pre", "options": 2},
-                {"name": "post", "options": 2},
-            ],
-        ),
-        SingleShift(
-            sites=("rota a",),
-            name="weekend a",
             workers_required=1,
             length=12.5,
             days=days[5:],
             balance_offset=2,
             constraint=[
-                {"name": "pre", "options": 2},
+                #{"name": "pre", "options": 2},
-                {"name": "post", "options": 2},
+                #{"name": "post", "options": 1},
             ],
             display_char="a",
+            #force_assign_with=["oncall"]
         ),
         SingleShift(
-            sites=("rota b", "rota c"),
+            sites=("rota b",),
             name="weekend b",
             workers_required=1,
             length=12.5,
@@ -113,22 +94,25 @@ def main(
                 "rota a",
                 "rota b",
             ),
-            name="evening",
+            name="twilight",
             workers_required=1,
             length=12.5,
             days=days[:5],
-            balance_offset=2,
+            balance_offset=5,
-            constraint=[],
+            #constraint=[
+            #    {"name": "pre", "options": 1},
+            #    {"name": "post", "options": 1},
+            #],
         ),
     )
-    Rota.allow_shifts_together(
+    #Rota.allow_shifts_together_for_all_workers(
-        
+    #        "oncall",
-            "oncall weekday",
+    #        "twilight",
-            "Oncall weekend",
+    #)
-            "weekend a",
+    #Rota.allow_shifts_together_for_all_workers(
-            "evening",
+    #        "oncall",
-        
+    #        "twilight",
-    )
+    #)
 
     # Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"])
 
@@ -136,8 +120,24 @@ def main(
     Rota.build_shifts()
     # Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names())
 
-    for w in range(1, 12):
+    for n in range(1, 12):
-        w = Worker(name=f"A{w:02d}", site="rota a", grade=1)
+        w = Worker(name=f"A{n:02d}", site="rota a", grade=1)
+        #w.prefer_multi_shift_together = 10
+        #w.allow_shifts_together("oncall", "twilight")
+        #w.allow_shifts_together("oncall", "weekend")
+
+        #w.add_hard_day_dependency("Fri", "Sat",  "Sun")
+
+        w.add_force_assign_with("twilight", "oncall")
+        w.add_force_assign_with("weekend", "oncall")
+        w.set_max_shifts_per_week("twilight", 1)
+
+        if n <= 2:
+            pass
+            w.add_hard_day_dependency("Fri", "Sat", "Sun")
+        else:
+            w.add_hard_day_dependency("Fri", "Sun")
+        ##    #w.add_hard_day_dependency("Sat", "Sun")
 
         print(w)
 
@@ -154,6 +154,9 @@ def main(
                 )
             ],
         )
+        #w.add_hard_day_exclusion("Sat", "Sun")
+        w.set_max_shifts_per_week("twilight", 1)
+        w.set_max_shifts_per_week("weekend b", 1)
 
         print(w)
 

rota/shifts.py

@@ -1385,6 +1385,7 @@ class RotaBuilder(object):
                         )
 
         # Most of our constraints apply per worker
+        # Worker constraint loop
         worker: Worker
         for worker in self.workers:
             console.rule(
@@ -1392,6 +1393,22 @@ class RotaBuilder(object):
             )
             logging.debug(f"Generate worker constraints: {worker.name}")
 
+
+            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():
+                    if shift_type not in self.get_shift_names():
+                        raise InvalidShift(
+                            f"Worker '{worker.name}' has max_shifts_per_week_by_shift_name referencing non-existent shift: {shift_type}"
+                        )
+                    for week in self.weeks:
+                        self.model.constraints.add(
+                            sum(
+                                self.model.works[worker.id, week, day, shift_type]
+                                for day in days
+                                if shift_type in self.get_shift_names_by_week_day(week, day)
+                            ) <= max_per_week
+                        )
+
             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:

rota/workers.py

@@ -117,6 +117,7 @@ class Worker(BaseModel):
     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"]}
+    max_shifts_per_week_by_shift_name: dict[str, int] = {}  # e.g. {"night": 2, "a": 3}
 
 
     assign_as_block_preferences: dict[str, float] = {}  # shift_name -> weight (positive = prefer block)
@@ -125,6 +126,8 @@ class Worker(BaseModel):
 
     weekend_shift_target_number: int = 0
 
+
+
     model_config = ConfigDict(
         extra="allow",
     )
@@ -422,3 +425,9 @@ class Worker(BaseModel):
         if isinstance(with_shifts, str):
             with_shifts = [with_shifts]
         self.force_assign_with[shift] = with_shifts
+
+    def set_max_shifts_per_week(self, shift_name: str, max_per_week: int):
+        """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

test/test_workers.py

@@ -1381,3 +1381,78 @@ def test_workers_double_shifts_on_weekends():
     Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
 
     assert Rota.results.solver.status == "ok"
+
+
+def test_max_shifts_per_week_by_shift_name_basic():
+    Rota = RotaBuilder(weeks_to_rota=2, start_date=datetime.date(2025, 6, 16))
+    w = Worker(name="A01", site="site1", grade=1)
+    w.max_shifts_per_week_by_shift_name = {"a": 1}
+    Rota.add_worker(w)
+    Rota.add_shifts(
+        SingleShift(
+            sites=("site1",),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=1,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    # Worker should not be assigned more than 1 "a" shift per week
+    for week in Rota.weeks:
+        count = sum(
+            Rota.model.works[w.id, week, day, "a"].value > 0.5
+            for day in days[:5]
+        )
+        assert count <= 1
+
+def test_max_shifts_per_week_by_shift_name_multiple_weeks():
+    Rota =generate_basic_rota(workers=4, weeks_to_rota=2)
+    for worker in Rota.get_workers():
+        worker.set_max_shifts_per_week("a", 1)
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=1,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    assert Rota.results.solver.status == "error"
+
+    Rota.add_worker(
+        Worker(name="extraworker", site="group1", grade=1, max_shifts_per_week_by_shift_name={"a": 1})
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    assert Rota.results.solver.status == "ok"
+
+    Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
+    # Worker should not be assigned more than 2 "a" shifts per week
+    for worker in Rota.get_workers():
+        for week in Rota.weeks:
+            count = sum(
+                Rota.model.works[worker.id, week, day, "a"].value > 0.5
+                for day in days[:5]
+            )
+            assert count <= 2
+
+
+def test_max_shifts_per_week_by_shift_name_invalid_shift():
+    Rota = RotaBuilder(weeks_to_rota=1, start_date=datetime.date(2025, 6, 16))
+    w = Worker(name="A03", site="site1", grade=1)
+    # "b" does not exist
+    w.max_shifts_per_week_by_shift_name = {"b": 1}
+    Rota.add_worker(w)
+    Rota.add_shifts(
+        SingleShift(
+            sites=("site1",),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=1,
+        ),
+    )
+    with pytest.raises(InvalidShift):
+        Rota.build_and_solve(options={"ratio": 0.0})