a few more fixes and improvements

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",
     )
@@ -421,4 +424,10 @@ class Worker(BaseModel):
             self.force_assign_with = {}
         if isinstance(with_shifts, str):
             with_shifts = [with_shifts]
-        self.force_assign_with[shift] = 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