feat: add group management for workers and shifts, including validation for overlapping groups

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)

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()}
@@ -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)
 
@@ -1610,13 +1623,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
+                    for worker in eligible_workers
-                    if worker.site in site_required
                 )
             )
 
@@ -1625,7 +1641,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 +1654,7 @@ class RotaBuilder(object):
                         # if not worker.locum
                     )
                 )
+
                 self.model.constraints.add(
                     0
                     == sum(
@@ -1647,14 +1664,17 @@ class RotaBuilder(object):
                     )
                 )
 
-            # And it is not assigned if the worker is from the wrong site
+            # And it is not assigned if the worker is ineligible
-            if [worker for worker in self.workers if worker.site not in site_required]:
+            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
+                        for worker in ineligible_workers
-                        if worker.site not in site_required
                     )
                 )
             # # Ensure shifts are only assigned by workers from the allowed sites
@@ -2647,14 +2667,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 +2691,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:
@@ -3825,7 +3845,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 +3860,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 +3874,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 +3896,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 +3911,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 +3925,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
                                 )
                             )
 
@@ -4181,7 +4201,7 @@ class RotaBuilder(object):
                 and shift.name in worker.assign_as_block_preferences
                 and (worker.id, week, shift.name)
                 in self.model.blocks_worker_shift_assigned
-                and worker.site in shift.sites
+                and self.is_worker_eligible_for_shift(worker, shift)
             )
 
             # Quadratic :(
@@ -4277,6 +4297,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 +4319,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 +4904,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 [shift for shift in self.shifts if self.is_worker_eligible_for_shift(worker, shift)]
-        return shifts
 
     def get_shifts_with_constraint(self, constraint) -> List[SingleShift]:
         return [shift for shift in self.shifts if shift.has_constraint(constraint)]
@@ -5066,8 +5096,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_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