feat: implement sub-block handling for shifts and update constraints in RotaBuilder

djangorota/rota/forms.py

@@ -46,6 +46,7 @@ class WorkerForm(forms.ModelForm):
     assign_as_block_preferences = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON mapping shift_name->weight", label="Block preferences (JSON)")
     shift_fte_overrides = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON mapping shift_name->fte", label="Shift FTE overrides (JSON)")
     exact_shifts = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON mapping shift_name->exact_count", label="Exact shifts (JSON)")
+    groups = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON list of group names (e.g. [\"group1\", \"group2\"])", label="Groups (JSON)")
     previous_shifts = forms.CharField(required=False, widget=forms.Textarea, help_text="Free JSON structure for previous shifts", label="Previous shifts (JSON)")
     shift_balance_extra = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON structure for extra shift-balance penalties", label="Shift balance extra (JSON)")
     allowed_multi_shift_sets = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON list of sets (e.g. [[\"a\",\"b\"], ...])", label="Allowed multi-shift sets (JSON)")
@@ -96,6 +97,7 @@ class WorkerForm(forms.ModelForm):
             "assign_as_block_preferences",
             "shift_fte_overrides",
             "exact_shifts",
+            "groups",
             "previous_shifts",
             "shift_balance_extra",
             "allowed_multi_shift_sets",
@@ -149,6 +151,7 @@ class WorkerForm(forms.ModelForm):
             "assign_as_block_preferences",
             "shift_fte_overrides",
             "exact_shifts",
+            "groups",
             "previous_shifts",
             "shift_balance_extra",
             "allowed_multi_shift_sets",
@@ -800,6 +803,16 @@ class ShiftForm(forms.Form):
         widget=forms.Textarea(attrs={"rows": 4, "class": "textarea"}),
         help_text="Optional JSON list of constraint objects for this shift (e.g. [{'name':'night','options':{}}])",
     )
+    include_groups = forms.CharField(
+        required=False,
+        help_text="Comma separated list of groups that can work this shift (leave empty for all)",
+        label="Include groups",
+    )
+    exclude_groups = forms.CharField(
+        required=False,
+        help_text="Comma separated list of groups that cannot work this shift",
+        label="Exclude groups",
+    )
 
     def clean_sites(self):
         val = self.cleaned_data["sites"]
@@ -841,3 +854,22 @@ class ShiftForm(forms.Form):
             raise
         except Exception as exc:
             raise forms.ValidationError(f"Invalid JSON for constraints: {exc}")
+
+    def clean_include_groups(self):
+        val = self.cleaned_data.get("include_groups") or ""
+        return [g.strip() for g in val.split(",") if g.strip()]
+
+    def clean_exclude_groups(self):
+        val = self.cleaned_data.get("exclude_groups") or ""
+        return [g.strip() for g in val.split(",") if g.strip()]
+
+    def clean(self):
+        cleaned_data = super().clean()
+        include = cleaned_data.get("include_groups") or []
+        exclude = cleaned_data.get("exclude_groups") or []
+        overlap = set(include).intersection(set(exclude))
+        if overlap:
+            raise forms.ValidationError(
+                f"Include groups and Exclude groups cannot overlap. Overlapping groups: {', '.join(overlap)}"
+            )
+        return cleaned_data

djangorota/rota/models.py

@@ -274,6 +274,17 @@ class Worker(models.Model):
 					else:
 						pyd_kwargs[k] = []
 
+				set_keys = ["groups"]
+				for k in set_keys:
+					v = pyd_kwargs.get(k, None)
+					try:
+						if v is None:
+							pyd_kwargs[k] = set()
+						else:
+							pyd_kwargs[k] = set(v)
+					except Exception:
+						pyd_kwargs[k] = set()
+
 				# allowed_multi_shift_sets is expected to be a list of sets; if the
 				# DB contains lists-of-lists convert inner lists to sets.
 				ams = pyd_kwargs.get("allowed_multi_shift_sets")

djangorota/rota/views.py

@@ -1690,6 +1690,8 @@ def shift_add(request, rota_id):
 				"workers_required": int(data["workers_required"]),
 				"assign_as_block": bool(data["assign_as_block"]),
 				"balance_offset": data.get("balance_offset"),
+				"include_groups": data.get("include_groups", []),
+				"exclude_groups": data.get("exclude_groups", []),
 				"constraints": data.get("constraints", []),
 			}
 
@@ -1738,6 +1740,9 @@ def shift_edit(request, rota_id, idx):
 				"days": data["days"],
 				"workers_required": int(data["workers_required"]),
 				"assign_as_block": bool(data["assign_as_block"]),
+				"balance_offset": data.get("balance_offset"),
+				"include_groups": data.get("include_groups", []),
+				"exclude_groups": data.get("exclude_groups", []),
 				"constraints": data.get("constraints", []),
 			}
 			current = rota.shifts or []
@@ -1766,6 +1771,8 @@ def shift_edit(request, rota_id, idx):
 		"workers_required": shift.get("workers_required"),
 		"assign_as_block": shift.get("assign_as_block", False),
 		"balance_offset": shift.get("balance_offset", None),
+		"include_groups": ",".join(shift.get("include_groups") or []),
+		"exclude_groups": ",".join(shift.get("exclude_groups") or []),
 		"constraints": json.dumps(shift.get("constraints", [])) if shift.get("constraints") is not None else "",
 	}
 	form = ShiftForm(initial=initial)

gen_proc.py

@@ -74,6 +74,7 @@ def main(
     Rota.constraint_options["max_weekend_frequency"] = 3
 
     Rota.constraint_options["max_days_per_week_block"] = [(8,3)]
+    Rota.constraint_options["maximum_allowed_shift_diff"] = 1
 
     # wr = [
     #    WorkerRequirement(
@@ -174,11 +175,11 @@ def main(
             length=12.5,
             days=days[4:],
             # balance_offset=3,
-            rota_on_nwds=True,
-            force_as_block=True,
+            #rota_on_nwds=True,
+            #force_as_block=True,
             # assign_as_block=True,
             constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
-            # force_as_block_unless_nwd=True
+            force_as_block_unless_nwd=[["Fri"], ["Sat", "Sun"]],
         ),
         SingleShift(
             sites=("torbay", "torbay twilights and weekends"),

rota_generator/shifts.py

@@ -346,7 +346,7 @@ def _pydantic_to_dict(obj):
     # Recursively convert pydantic models, lists and dicts to JSON-serializable structures
     if obj is None:
         return None
-    if isinstance(obj, list):
+    if isinstance(obj, (list, set, tuple)):
         return [_pydantic_to_dict(x) for x in obj]
     if isinstance(obj, dict):
         return {k: _pydantic_to_dict(v) for k, v in obj.items()}
@@ -417,7 +417,7 @@ class SingleShift(BaseModel):
     rota_on_nwds: bool = False
     assign_as_block: bool = False
     force_as_block: bool = False
-    force_as_block_unless_nwd: bool = False
+    force_as_block_unless_nwd: bool | List[List[str]] = False
     hard_constrain_shift: bool = True
     bank_holidays_only: bool = False
     #constraint: list[ShiftConstraint] = []
@@ -433,6 +433,17 @@ class SingleShift(BaseModel):
     pair_proxy: bool = False
     display_char: str | None = None  # Add this line
     force_assign_with: list[str] = []  # List of shift names to force assign with this shift on the same day
+    include_groups: set[str] = set()
+    exclude_groups: set[str] = set()
+
+    @model_validator(mode="after")
+    def check_groups_non_overlapping(self) -> "SingleShift":
+        intersection = self.include_groups.intersection(self.exclude_groups)
+        if intersection:
+            raise ValueError(
+                f"include_groups and exclude_groups cannot overlap. Intersection: {intersection}"
+            )
+        return self
 
     model_config = ConfigDict(
         extra="allow",
@@ -855,6 +866,8 @@ class RotaBuilder(object):
                 options["time_limit"] = options.pop("seconds")
             if "ratio" in options:
                 options["mip_rel_gap"] = options.pop("ratio")
+            if "threads" in options:
+                options.pop("threads")
         else:
             self.opt = SolverFactory(solver)
 
@@ -1037,20 +1050,28 @@ class RotaBuilder(object):
             initialize=0,
         )
 
+        blocks_worker_keys = []
+        for worker, week, shift_name in self.worker_week_shifts_to_assign_as_blocks():
+            shift = self.get_shift_by_name(shift_name)
+            sub_blocks = self.get_shift_sub_blocks(shift)
+            for sb_idx in range(len(sub_blocks)):
+                blocks_worker_keys.append((worker.id, week, shift_name, sb_idx))
+
         self.model.blocks_worker_shift_assigned = Var(
-            (
-                (worker.id, week, shift)
-                for worker, week, shift in self.worker_week_shifts_to_assign_as_blocks()
-            ),
+            blocks_worker_keys,
             within=Binary,
             initialize=0,
         )
 
+        blocks_assigned_keys = []
+        for week, shift_name in self.week_shifts_to_assign_as_blocks():
+            shift = self.get_shift_by_name(shift_name)
+            sub_blocks = self.get_shift_sub_blocks(shift)
+            for sb_idx in range(len(sub_blocks)):
+                blocks_assigned_keys.append((week, shift_name, sb_idx))
+
         self.model.blocks_assigned = Var(
-            (
-                (week, shift)
-                for week, shift in self.week_shifts_to_assign_as_blocks()
-            ),
+            blocks_assigned_keys,
             within=NonNegativeIntegers,
             initialize=0,
         )
@@ -1610,13 +1631,16 @@ class RotaBuilder(object):
             workers_required,
             site_required,
         ) in self.get_required_workers_and_site_combinations():
-            # print(week, day, shift, workers_required, site_required)
+            shift_obj = self.get_shift_by_name(shift)
+            eligible_workers = [
+                worker for worker in self.workers
+                if self.is_worker_eligible_for_shift(worker, shift_obj)
+            ]
             self.model.constraints.add(
                 workers_required
                 == sum(
                     self.model.works[worker.id, week, day, shift]
-                    for worker in self.workers
-                    if worker.site in site_required
+                    for worker in eligible_workers
                 )
             )
 
@@ -1625,7 +1649,7 @@ class RotaBuilder(object):
                     self.model.locum_required[week, day, shift]
                     == sum(
                         self.model.works[worker.id, week, day, shift]
-                        for worker in self.workers
+                        for worker in eligible_workers
                         if worker.locum
                     )
                 )
@@ -1638,6 +1662,7 @@ class RotaBuilder(object):
                         # if not worker.locum
                     )
                 )
+
                 self.model.constraints.add(
                     0
                     == sum(
@@ -1647,14 +1672,17 @@ class RotaBuilder(object):
                     )
                 )
 
-            # And it is not assigned if the worker is from the wrong site
-            if [worker for worker in self.workers if worker.site not in site_required]:
+            # And it is not assigned if the worker is ineligible
+            ineligible_workers = [
+                worker for worker in self.workers
+                if not self.is_worker_eligible_for_shift(worker, shift_obj)
+            ]
+            if ineligible_workers:
                 self.model.constraints.add(
                     0
                     == sum(
                         self.model.works[worker.id, week, day, shift]
-                        for worker in self.workers
-                        if worker.site not in site_required
+                        for worker in ineligible_workers
                     )
                 )
             # # Ensure shifts are only assigned by workers from the allowed sites
@@ -1880,99 +1908,103 @@ class RotaBuilder(object):
 
         for week, shift_name in self.week_shifts_to_assign_as_blocks():
                 shift = self.get_shift_by_name(shift_name)
+                sub_blocks = self.get_shift_sub_blocks(shift)
 
-                for worker in self.get_workers_for_shift(shift):
-                    if (worker.id, week, shift.name) in self.model.blocks_worker_shift_assigned:
-                        try:
-                            self.model.constraints.add(
-                                8
-                                * self.model.blocks_worker_shift_assigned[
-                                    worker.id, week, shift_name
-                                ]
-                                >= sum(
-                                    self.model.works[worker.id, week, day, shift_name]
-                                    for day in shift.days
+                for sb_idx, sb in enumerate(sub_blocks):
+                    for worker in self.get_workers_for_shift(shift):
+                        if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned:
+                            try:
+                                self.model.constraints.add(
+                                    8
+                                    * self.model.blocks_worker_shift_assigned[
+                                        worker.id, week, shift_name, sb_idx
+                                    ]
+                                    >= sum(
+                                        self.model.works[worker.id, week, day, shift_name]
+                                        for day in sb
+                                    )
                                 )
-                            )
-                        except KeyError:
-                            pass
-                self.model.constraints.add(
-                    self.model.blocks_assigned[week, shift_name]
-                    == sum(
-                        self.model.blocks_worker_shift_assigned[
-                            worker.id, week, shift_name
-                        ]
-                        for worker in self.workers
-                        if (worker.id, week, shift.name) in self.model.blocks_worker_shift_assigned
+                            except KeyError:
+                                pass
+
+                    self.model.constraints.add(
+                        self.model.blocks_assigned[week, shift_name, sb_idx]
+                        == sum(
+                            self.model.blocks_worker_shift_assigned[
+                                worker.id, week, shift_name, sb_idx
+                            ]
+                            for worker in self.workers
+                            if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
+                        )
                     )
-                )
 
-                workers_required = shift.get_worker_requirement_by_date(
-                    self.get_week_start_date(week)
-                )
+                    workers_required = shift.get_worker_requirement_by_date(
+                        self.get_week_start_date(week)
+                    )
 
-                if shift.force_as_block_unless_nwd:
-                    workers = self.get_workers_for_shift(shift)
-                    # Get workers who have a nwd on the shift
-                    nwd_workers = []
-                    full_workers = []
+                    if shift.force_as_block_unless_nwd:
+                        workers = self.get_workers_for_shift(shift)
+                        # Get workers who have a nwd on the sub-block
+                        nwd_workers = []
+                        full_workers = []
 
-                    for w in workers:
-                        if w.non_working_day_list:
-                            l = []
-                            for (
-                                nwd,
-                                start_nwd_date,
-                                end_nwd_date,
-                            ) in w.non_working_day_list:
-                                if nwd in shift.days:
-                                    if start_nwd_date > self.get_week_start_date(week):
-                                        continue
-                                    if end_nwd_date < self.get_week_start_date(week):
-                                        continue
-                                    l.append(w)
+                        for w in workers:
+                            if w.non_working_day_list:
+                                l = []
+                                for (
+                                    nwd,
+                                    start_nwd_date,
+                                    end_nwd_date,
+                                ) in w.non_working_day_list:
+                                    if nwd in sb:
+                                        if start_nwd_date > self.get_week_start_date(week):
+                                            continue
+                                        if end_nwd_date < self.get_week_start_date(week):
+                                            continue
+                                        l.append(w)
 
-                            if l:
-                                nwd_workers.append(w)
+                                if l:
+                                    nwd_workers.append(w)
+                                else:
+                                    full_workers.append(w)
                             else:
                                 full_workers.append(w)
 
+                        if not nwd_workers:  # Just do the usual
+                            self.model.constraints.add(
+                                sum(
+                                    self.model.blocks_worker_shift_assigned[
+                                        worker.id, week, shift.name, sb_idx
+                                    ]
+                                    for worker in self.workers
+                                    if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
+                                )
+                                <= workers_required
+                            )
                         else:
-                            full_workers.append(w)
-
-                    if not nwd_workers:  # Just do the usual
-                        self.model.constraints.add(
-                            sum(
-                                self.model.blocks_worker_shift_assigned[
-                                    worker.id, week, shift.name
-                                ]
-                                for worker in self.workers
+                            self.model.constraints.add(
+                                sum(
+                                    self.model.blocks_worker_shift_assigned[
+                                        worker.id, week, shift.name, sb_idx
+                                    ]
+                                    for worker in full_workers
+                                    if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
+                                )
+                                <= workers_required + 1
                             )
-                            <= workers_required
-                        )
 
-                    else:
-                        self.model.constraints.add(
-                            sum(
-                                self.model.blocks_worker_shift_assigned[
-                                    worker.id, week, shift.name
-                                ]
-                                for worker in full_workers
+                    elif shift.force_as_block:
+                        if workers_required:
+                            self.model.constraints.add(
+                                sum(
+                                    self.model.blocks_worker_shift_assigned[
+                                        worker.id, week, shift.name, sb_idx
+                                    ]
+                                    for worker in self.workers
+                                    if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
+                                )
+                                <= workers_required
                             )
-                            <= workers_required + 1
-                        )
-
-                elif shift.force_as_block:
-                    if workers_required:
-                        self.model.constraints.add(
-                            sum(
-                                self.model.blocks_worker_shift_assigned[
-                                    worker.id, week, shift.name
-                                ]
-                                for worker in self.workers
-                            )
-                            <= workers_required
-                        )
 
         # Most of our constraints apply per worker
         # Worker constraint loop (worker loop)
@@ -2647,14 +2679,14 @@ class RotaBuilder(object):
                 self.get_shifts(), description="Generate shift balance constraints"
             ):
                 if (
-                    worker.site in shift.sites
+                    self.is_worker_eligible_for_shift(worker, shift)
                 ):  # Each site specfies which sites self.workers can fullfill
                     total_shifts = self.shift_worker_counts[shift.name]
 
                     # Find workers who have exact shifts for this shift type
                     exact_workers = [
                         w for w in self.workers 
-                        if w.site in shift.sites and shift.name in getattr(w, "exact_shifts", {})
+                        if self.is_worker_eligible_for_shift(w, shift) and shift.name in getattr(w, "exact_shifts", {})
                     ]
                     exact_shifts_sum = sum(w.exact_shifts[shift.name] for w in exact_workers)
 
@@ -2671,7 +2703,7 @@ class RotaBuilder(object):
                         full_time_equivalent_joined = sum(
                             w.get_fte(shift=shift.name)
                             for w in self.workers
-                            if w.site in shift.sites and w not in exact_workers
+                            if self.is_worker_eligible_for_shift(w, shift) and w not in exact_workers
                         )
 
                         if not full_time_equivalent_joined:
@@ -2768,18 +2800,19 @@ class RotaBuilder(object):
                             "target_shifts": target_shifts,
                         }
 
-                        self.model.constraints.add(
-                            inequality(
-                                min_shifts,
-                                sum(
-                                    self.model.works[worker.id, week, day, shift.name]
-                                    for week, day in self.get_week_day_combinations_for_shift(
-                                        shift
-                                    )
-                                ),
-                                max_shifts,
+                        if self.get_week_day_combinations_for_shift(shift):
+                            self.model.constraints.add(
+                                inequality(
+                                    min_shifts,
+                                    sum(
+                                        self.model.works[worker.id, week, day, shift.name]
+                                        for week, day in self.get_week_day_combinations_for_shift(
+                                            shift
+                                        )
+                                    ),
+                                    max_shifts,
+                                )
                             )
-                        )
                 else:
                     # Make sure shifts aren't assigned to those from other sites
                     # Not needed if there are no shifts
@@ -3825,7 +3858,7 @@ class RotaBuilder(object):
                             continue
                         if day in constraint_shift.days:
                             # Only apply if worker can work this shift
-                            if worker.site not in constraint_shift.sites:
+                            if not self.is_worker_eligible_for_shift(worker, constraint_shift):
                                 continue
                             try:
                                 works = self.model.works[
@@ -3840,7 +3873,7 @@ class RotaBuilder(object):
                                 pre_date = None
                             if pre_date is not None and pre_date in getattr(constraint, "exclude_dates", []):
                                 continue
-
+ 
                             self.model.constraints.add(
                                 1
                                 >= works
@@ -3854,7 +3887,7 @@ class RotaBuilder(object):
                                     )
                                     if shiftname not in ignore_shifts
                                     for w in workers
-                                    if w.site in constraint_shift.sites  # Only workers who can work this shift
+                                    if self.is_worker_eligible_for_shift(w, constraint_shift)  # Only workers who can work this shift
                                 )
                             )
 
@@ -3876,7 +3909,7 @@ class RotaBuilder(object):
                             continue
                         if day in constraint_shift.days:
                             # Only apply if worker can work this shift
-                            if worker.site not in constraint_shift.sites:
+                            if not self.is_worker_eligible_for_shift(worker, constraint_shift):
                                 continue
                             try:
                                 works = self.model.works[
@@ -3891,7 +3924,7 @@ class RotaBuilder(object):
                                 post_date = None
                             if post_date is not None and post_date in getattr(constraint, "exclude_dates", []):
                                 continue
-
+ 
                             self.model.constraints.add(
                                 1
                                 >= works
@@ -3905,7 +3938,7 @@ class RotaBuilder(object):
                                     )
                                     if shiftname not in ignore_shifts
                                     for w in workers
-                                    if w.site in constraint_shift.sites  # Only workers who can work this shift
+                                    if self.is_worker_eligible_for_shift(w, constraint_shift)  # Only workers who can work this shift
                                 )
                             )
 
@@ -4164,24 +4197,27 @@ class RotaBuilder(object):
             if self.constraint_options_model.balance_blocks:
                 blocks_balancing = sum(
                     block_shift_balancing_constant
-                    * self.model.blocks_assigned[week, shift]
+                    * self.model.blocks_assigned[week, shift, sb_idx]
                     for week in self.weeks
                     for shift in self.shifts_to_assign_as_blocks()
+                    for sb_idx in range(len(self.get_shift_sub_blocks(self.get_shift_by_name(shift))))
+                    if (week, shift, sb_idx) in self.model.blocks_assigned
                 )
             else:
                 blocks_balancing = 0
 
             prefer_block_expr = sum(
                 worker.assign_as_block_preferences.get(shift.name, 0)
-                * self.model.blocks_worker_shift_assigned[worker.id, week, shift.name]
+                * self.model.blocks_worker_shift_assigned[worker.id, week, shift.name, sb_idx]
                 for worker in self.workers
                 for week in self.weeks
                 for shift in self.shifts
                 if hasattr(worker, "assign_as_block_preferences")
                 and shift.name in worker.assign_as_block_preferences
-                and (worker.id, week, shift.name)
+                and self.is_worker_eligible_for_shift(worker, shift)
+                for sb_idx in range(len(self.get_shift_sub_blocks(shift)))
+                if (worker.id, week, shift.name, sb_idx)
                 in self.model.blocks_worker_shift_assigned
-                and worker.site in shift.sites
             )
 
             # Quadratic :(
@@ -4277,6 +4313,11 @@ class RotaBuilder(object):
                             "Invalid exact shift site",
                             f"Worker {worker.name} requested exact shifts for shift {s_name} but their site {worker.site} is not in the shift sites {s.sites}",
                         )
+                    elif not self.is_worker_eligible_for_shift(worker, s):
+                        self.add_warning(
+                            "Invalid exact shift group",
+                            f"Worker {worker.name} requested exact shifts for shift {s_name} but is not eligible for it",
+                        )
                 except KeyError:
                     self.add_warning(
                         "Invalid exact shift",
@@ -4294,7 +4335,13 @@ class RotaBuilder(object):
                 self.add_warning("Worker/duplicate name", message)
             self.workers_name_map[worker.name] = worker
 
-            if worker.site not in self.sites:
+            worker_has_valid_shift = False
+            for shift in self.shifts:
+                if self.is_worker_eligible_for_shift(worker, shift):
+                    worker_has_valid_shift = True
+                    break
+
+            if not worker_has_valid_shift:
                 message = f"Worker with name '{worker.name}' ({worker.id}) has no valid shifts (site: {worker.site})"
                 logger.warning(message)
                 self.add_warning("Worker/no valid shifts", message)
@@ -4873,8 +4920,7 @@ class RotaBuilder(object):
 
     def get_shifts_for_worker(self, worker_id):
         worker = self.get_worker_by_id(worker_id)
-        shifts = [shift for shift in self.shifts if worker.site in shift.sites]
-        return shifts
+        return [shift for shift in self.shifts if self.is_worker_eligible_for_shift(worker, shift)]
 
     def get_shifts_with_constraint(self, constraint) -> List[SingleShift]:
         return [shift for shift in self.shifts if shift.has_constraint(constraint)]
@@ -5015,6 +5061,27 @@ class RotaBuilder(object):
             )
         ])
 
+    def get_shift_sub_blocks(self, shift: SingleShift) -> List[List[str]]:
+        if getattr(shift, "force_as_block_unless_nwd", None):
+            val = shift.force_as_block_unless_nwd
+            if isinstance(val, (list, tuple)) and all(isinstance(x, (list, tuple, set)) for x in val):
+                return [list(x) for x in val]
+            return [list(shift.days)]
+        
+        # Check if it is a block shift globally or per-worker
+        is_block = (
+            shift.assign_as_block
+            or shift.force_as_block
+            or any(
+                hasattr(worker, "assign_as_block_preferences")
+                and getattr(worker, "assign_as_block_preferences", {}).get(shift.name, 0) != 0
+                for worker in self.workers
+            )
+        )
+        if is_block:
+            return [list(shift.days)]
+        return []
+
     def get_all_locum_availability(self):
         return self.locum_availability_map
 
@@ -5066,8 +5133,46 @@ class RotaBuilder(object):
 
         return group_workers
 
+    def is_worker_eligible_for_shift(self, worker: Worker, shift: SingleShift) -> bool:
+        """Returns True if the worker is eligible for the shift based on site and group options."""
+        # 1. Group checks
+        worker_groups = getattr(worker, "groups", set())
+        if not isinstance(worker_groups, (set, list, tuple)):
+            worker_groups = {worker_groups} if worker_groups else set()
+        worker_groups_set = set(worker_groups)
+
+        # exclude_groups check: worker MUST NOT belong to any group in exclude_groups
+        if getattr(shift, "exclude_groups", None):
+            exclude_set = set(shift.exclude_groups)
+            if worker_groups_set.intersection(exclude_set):
+                return False
+
+        # Additive Site or Include Group check
+        has_site = worker.site in shift.sites
+        has_include_group = False
+        if getattr(shift, "include_groups", None):
+            include_set = set(shift.include_groups)
+            if worker_groups_set.intersection(include_set):
+                has_include_group = True
+
+        if getattr(shift, "include_groups", None):
+            if not (has_site or has_include_group):
+                return False
+        else:
+            if not has_site:
+                return False
+
+        return True
+
+    def get_workers_in_group(self, group_name: str) -> List[Worker]:
+        """Returns a list of all workers belonging to the specified group."""
+        return [
+            w for w in self.workers
+            if group_name in getattr(w, "groups", set())
+        ]
+
     def get_workers_for_shift(self, shift: SingleShift) -> List[Worker]:
-        return [worker for worker in self.workers if worker.site in shift.sites]
+        return [worker for worker in self.workers if self.is_worker_eligible_for_shift(worker, shift)]
 
     def get_workers_total_fte(self) -> float:
         """Does not take into account shift adjusted ftes"""

rota_generator/workers.py

@@ -214,6 +214,7 @@ class Worker(BaseModel):
     shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist
 
     exact_shifts: dict[str, int] = {} # Map shift_name to exact number of shifts
+    groups: set[str] = set() # Groups that the worker belongs to
 
     weekend_shift_target_number: float = 0.0
     avoid_shifts_on_dates: list[AvoidShiftOnDates] = []

test/test_nwds.py

@@ -163,7 +163,7 @@ def test_nwd_force_as_block_force_split():
     for worker in Rota.workers:
         shifts = Rota.get_worker_shift_list(worker)
         shifts_string = "".join([i if i != "" else "-" for i in shifts])
-        for week in weeks_from_list(shifts_string[7 * 5 :]):
+        for week in weeks_from_list(shifts_string[7 * 6 :]):
             assert week in ("-----ww", "dddddww")
         if worker.name == "worker1":
             assert shifts_string[: 7 * 5].count("ww-") == 4
@@ -213,4 +213,61 @@ def test_nwd_testing():
             if worker.name == "worker1":
                 assert week == "dddddww"
             else:
-                assert week == "-----ww"
+                assert week == "-----ww"
+
+
+def test_force_as_block_unless_nwd_with_subblocks():
+    # Test that a Fri/Sat/Sun shift can be split into Fri and Sat/Sun blocks
+    # if a worker does not work on Fridays.
+    Rota = setup_rota(weeks_to_rota=2)
+    start_date = Rota.start_date
+    
+    # worker1 does not work on Fridays (Fri is NWD)
+    worker1 = Worker(
+        name="worker1", site="group1", grade=1,
+        nwds=[{"day": "Fri", "start_date": start_date, "end_date": start_date + datetime.timedelta(weeks=2)}],
+    )
+    # worker2 works normally
+    worker2 = Worker(name="worker2", site="group1", grade=1)
+    
+    Rota.add_workers((worker1, worker2))
+    
+    # Fri/Sat/Sun shift, requiring 1 worker, forced as block unless NWD
+    # We specify sub-blocks: [["Fri"], ["Sat", "Sun"]]
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="weekend", length=12.5,
+            days=["Fri", "Sat", "Sun"],
+            workers_required=1,
+            force_as_block_unless_nwd=[["Fri"], ["Sat", "Sun"]],
+        ),
+    )
+    
+    # We remove no valid shifts warning
+    Rota.terminate_on_warning.remove("Worker/no valid shifts")
+    
+    Rota.build_and_solve()
+    assert Rota.results.solver.status == "ok"
+    
+    # Check assignments:
+    # Since worker1 has NWD on Friday:
+    # - worker1 cannot work Friday.
+    # - But worker1 can work Sat and Sun as a block!
+    # So on Saturday and Sunday, worker1 should be assigned.
+    # On Friday, worker2 should be assigned.
+    for week in Rota.weeks:
+        w1_fri = Rota.model.works[worker1.id, week, "Fri", "weekend"].value
+        w1_sat = Rota.model.works[worker1.id, week, "Sat", "weekend"].value
+        w1_sun = Rota.model.works[worker1.id, week, "Sun", "weekend"].value
+        
+        w2_fri = Rota.model.works[worker2.id, week, "Fri", "weekend"].value
+        w2_sat = Rota.model.works[worker2.id, week, "Sat", "weekend"].value
+        w2_sun = Rota.model.works[worker2.id, week, "Sun", "weekend"].value
+        
+        assert w1_fri == 0
+        
+        # Either worker1 works Sat/Sun and worker2 works Fri (split block)
+        # OR worker2 works Fri/Sat/Sun (full block) and worker1 works nothing
+        is_split = (w1_sat == 1 and w1_sun == 1 and w2_fri == 1 and w2_sat == 0 and w2_sun == 0)
+        is_full_w2 = (w1_sat == 0 and w1_sun == 0 and w2_fri == 1 and w2_sat == 1 and w2_sun == 1)
+        assert is_split or is_full_w2

test/test_workers.py

@@ -1,5 +1,6 @@
 import pytest
 from rota_generator.shifts import InvalidShift, MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days
+from pydantic import ValidationError
 
 import datetime
 from rota_generator.workers import NonWorkingDays, NotAvailableToWork, Worker, generate_not_available_to_works
@@ -1868,4 +1869,119 @@ def test_html_export_worker_details_and_unavailable_reason():
     html = Rota.get_worker_timetable_html()
     
     assert 'data-exact-shifts=\'{"a": 1}\'' in html
-    assert 'title=\' (2022-03-07) / START DATE: 2022-03-14\'' in html
+    assert 'title=\' (2022-03-07) / START DATE: 2022-03-14\'' in html
+
+
+def test_worker_groups_selector():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group1", grade=1, fte=100, groups={"groupA", "groupB"}),
+        Worker(name="worker2", site="group1", grade=1, fte=100, groups={"groupB", "groupC"}),
+    ])
+    
+    workers_A = Rota.get_workers_in_group("groupA")
+    workers_B = Rota.get_workers_in_group("groupB")
+    workers_C = Rota.get_workers_in_group("groupC")
+    workers_D = Rota.get_workers_in_group("groupD")
+    
+    assert [w.name for w in workers_A] == ["worker1"]
+    assert set(w.name for w in workers_B) == {"worker1", "worker2"}
+    assert [w.name for w in workers_C] == ["worker2"]
+    assert workers_D == []
+
+
+def test_shift_groups_overlapping_validation():
+    with pytest.raises(ValidationError) as exc_info:
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            include_groups={"pool1", "pool2"},
+            exclude_groups={"pool2", "pool3"},
+        )
+    assert "include_groups and exclude_groups cannot overlap" in str(exc_info.value)
+
+
+def test_shift_group_availability_inclusion():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group1", grade=1, fte=100),
+        Worker(name="worker2", site="group2", grade=1, fte=100, groups={"senior"}),
+        Worker(name="worker3", site="group2", grade=1, fte=100, groups={"junior"}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            include_groups={"senior"},
+        ),
+    )
+    Rota.terminate_on_warning.remove("Worker/no valid shifts")
+    Rota.build_and_solve()
+    Rota.export_rota_to_html("test_shift_group_availability_inclusion", folder="tests")
+    assert Rota.results.solver.status == "ok"
+    
+    worker_shift_counts = {w.name: 0 for w in Rota.get_workers()}
+    for week, day, shiftname in Rota.get_all_shiftname_combinations():
+        for worker in Rota.get_workers():
+            val = Rota.model.works[worker.id, week, day, shiftname].value
+            if val is not None and val > 0.5:
+                worker_shift_counts[worker.name] += 1
+                
+    assert worker_shift_counts["worker1"] in (1, 2)
+    assert worker_shift_counts["worker2"] in (1, 2)
+    assert worker_shift_counts["worker3"] == 0
+
+
+def test_shift_group_availability_exclusion():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group1", grade=1, fte=100, groups={"senior"}),
+        Worker(name="worker2", site="group1", grade=1, fte=100, groups={"junior"}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            exclude_groups={"junior"},
+        ),
+    )
+    Rota.terminate_on_warning.remove("Worker/no valid shifts")
+    Rota.build_and_solve()
+    assert Rota.results.solver.status == "ok"
+    
+    worker_shift_counts = {w.name: 0 for w in Rota.get_workers()}
+    for week, day, shiftname in Rota.get_all_shiftname_combinations():
+        for worker in Rota.get_workers():
+            val = Rota.model.works[worker.id, week, day, shiftname].value
+            if val is not None and val > 0.5:
+                worker_shift_counts[worker.name] += 1
+                
+    assert worker_shift_counts["worker1"] == 2
+    assert worker_shift_counts["worker2"] == 0