a few more fixes and improvements

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