From b5525bd3fb98535e3eefaa35820c222f117fdc39 Mon Sep 17 00:00:00 2001 From: Ross Date: Thu, 9 Jul 2026 22:14:30 +0100 Subject: [PATCH] feat: add group management for workers and shifts, including validation for overlapping groups --- djangorota/rota/forms.py | 32 +++++++++++ djangorota/rota/models.py | 11 ++++ djangorota/rota/views.py | 7 +++ rota_generator/shifts.py | 114 ++++++++++++++++++++++++++++-------- rota_generator/workers.py | 1 + test/test_workers.py | 118 +++++++++++++++++++++++++++++++++++++- 6 files changed, 259 insertions(+), 24 deletions(-) diff --git a/djangorota/rota/forms.py b/djangorota/rota/forms.py index 4dd4e4c..635f97c 100644 --- a/djangorota/rota/forms.py +++ b/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 diff --git a/djangorota/rota/models.py b/djangorota/rota/models.py index 5dcb07e..7102d79 100644 --- a/djangorota/rota/models.py +++ b/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") diff --git a/djangorota/rota/views.py b/djangorota/rota/views.py index ffa137f..a05a669 100644 --- a/djangorota/rota/views.py +++ b/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) diff --git a/rota_generator/shifts.py b/rota_generator/shifts.py index 2a571f2..da5c53f 100644 --- a/rota_generator/shifts.py +++ b/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 - if worker.site in site_required + for worker in eligible_workers ) ) @@ -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 - 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 @@ -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 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)] @@ -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""" diff --git a/rota_generator/workers.py b/rota_generator/workers.py index e37458e..4ab29cb 100644 --- a/rota_generator/workers.py +++ b/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] = [] diff --git a/test/test_workers.py b/test/test_workers.py index 54f66e2..732185d 100644 --- a/test/test_workers.py +++ b/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 \ No newline at end of file + 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 \ No newline at end of file