numberous fixes and improvements

rota/shifts.py

@@ -145,6 +145,7 @@ class SingleShift(BaseModel):
     end_date: datetime.date | None = None
     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
 
     model_config = ConfigDict(
         extra="allow",
@@ -1391,17 +1392,39 @@ class RotaBuilder(object):
             )
             logging.debug(f"Generate worker constraints: {worker.name}")
 
+            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:
+                    for other_shift_name in shift.force_assign_with:
+                        # Only add constraint if the other shift is available on this day
+                        if other_shift_name in self.get_shift_names_by_week_day(week, day):
+                            for worker in self.get_workers_for_shift(shift):
+                                self.model.constraints.add(
+                                    self.model.works[worker.id, week, day, shift.name]
+                                    <= self.model.works[worker.id, week, day, other_shift_name]
+                                )
+
             # Add hard shift dependencies for this worker
             if hasattr(worker, "hard_day_dependencies"):
-                for from_day, to_day in worker.hard_day_dependencies.items():
+                for day_group in worker.hard_day_dependencies:
                     for week in self.weeks:
-                        for shift in self.get_shift_names_by_week_day(week, from_day):
+                        assigned_any_shift = {}
-                            # Only apply if the shift also exists on the to_day
+                        for day in day_group:
-                            if shift in self.get_shift_names_by_week_day(week, to_day):
+                            shifts_today = self.get_shift_names_by_week_day(week, day)
-                                self.model.constraints.add(
+                            var_name = f"hard_day_dep_{worker.id}_{week}_{day}"
-                                    self.model.works[worker.id, week, from_day, shift]
+                            if not hasattr(self.model, "hard_day_dep_vars"):
-                                    <= self.model.works[worker.id, week, to_day, shift]
+                                self.model.hard_day_dep_vars = {}
-                                )
+                            if (worker.id, week, day) not in self.model.hard_day_dep_vars:
+                                self.model.hard_day_dep_vars[(worker.id, week, day)] = Var(within=Binary)
+                                setattr(self.model, var_name, self.model.hard_day_dep_vars[(worker.id, week, day)])
+                            assigned_any_shift[day] = self.model.hard_day_dep_vars[(worker.id, week, day)]
+                            total_assigned = sum(self.model.works[worker.id, week, day, shift] for shift in shifts_today)
+                            self.model.constraints.add(total_assigned >= assigned_any_shift[day])
+                            self.model.constraints.add(total_assigned <= len(shifts_today) * assigned_any_shift[day])
+                        # All-or-none: all days in the group must have the same assignment status
+                        vals = list(assigned_any_shift.values())
+                        for i in range(1, len(vals)):
+                            self.model.constraints.add(vals[i] == vals[0])
 
             # Add hard day exclusions for this worker
             if hasattr(worker, "hard_day_exclusions"):
@@ -1448,6 +1471,17 @@ class RotaBuilder(object):
                             - len(allowed_set)
                             + 1
                         )
+                    shifts_today = self.get_shift_names_by_week_day(week, day)
+                for shift_name in shifts_today:
+                    if hasattr(worker, "force_assign_with") and shift_name in worker.force_assign_with:
+                        for other_shift in worker.force_assign_with[shift_name]:
+                            if other_shift in shifts_today:
+                                self.model.constraints.add(
+                                    self.model.works[worker.id, week, day, shift_name]
+                                    <= self.model.works[worker.id, week, day, other_shift]
+                                )
+
+
 
             if self.get_workers_who_require_locums():
                 for week, day, shift in self.get_all_shiftclass_combinations():
@@ -2215,17 +2249,20 @@ class RotaBuilder(object):
                 #         for day in self.days[5:]
                 #         for shift in self.get_shifts()) / 2)
                 # if self.constraint_options["balance_weekends"]:
-                self.model.constraints.add(
+
-                    self.model.works_weekend[worker.id, week]
+
-                    >= sum(
+                # Try disabling tihs to allow more than 2 shifts ot be assigned on a weekend
-                        self.model.works[worker.id, week, day, shiftname]
+                #self.model.constraints.add(
-                        for w, day, shiftname in self.get_all_shiftname_combinations(
+                #    self.model.works_weekend[worker.id, week]
-                            week=week
+                #    >= sum(
-                        )
+                #        self.model.works[worker.id, week, day, shiftname]
-                        if day in self.days[5:]
+                #        for w, day, shiftname in self.get_all_shiftname_combinations(
-                    )
+                #            week=week
-                    / 2
+                #        )
-                )
+                #        if day in self.days[5:]
+                #    )
+                #    / 2
+                #)
                 self.model.constraints.add(
                     self.model.works_weekend[worker.id, week]
                     <= sum(
@@ -2492,11 +2529,24 @@ class RotaBuilder(object):
                     worker_allowed_sets = getattr(
                         worker, "allowed_multi_shift_sets", []
                     )
+                    force_assign_sets = []
+                    for shift_name in shifts_today:
+                        if hasattr(worker, "force_assign_with") and shift_name in worker.force_assign_with:
+                            force_set = set([shift_name] + list(worker.force_assign_with[shift_name]))
+                            if force_set.issubset(shifts_today):
+                                force_assign_sets.append(force_set)
+                        shift_obj = self.get_shift_by_name(shift_name)
+                        if getattr(shift_obj, "force_assign_with", []):
+                            force_set = set([shift_name] + list(shift_obj.force_assign_with))
+                            if force_set.issubset(shifts_today):
+                                force_assign_sets.append(force_set)
+
+                    # If no allowed sets, we can only assign one shift
                     allowed_sets_today = [
                         allowed_set
                         for allowed_set in worker_allowed_sets
                         if allowed_set.issubset(shifts_today)
-                    ]
+                    ] + force_assign_sets
 
                     if allowed_sets_today:
                         # Forbid any multi-shift assignment that includes shifts from different allowed sets
@@ -3182,6 +3232,24 @@ class RotaBuilder(object):
             else:
                 s.end_date = self.rota_end_date
 
+
+            if hasattr(s, "force_assign_with") and s.force_assign_with:
+                missing = set(s.force_assign_with) - set(shift.name for shift in self.shifts)
+                if missing:
+                    raise InvalidShift(
+                        f"Shift '{s.name}' has force_assign_with referencing non-existent shift(s): {missing}"
+                    )
+
+            all_shift_names = set(shift.name for shift in self.shifts)
+            for worker in self.workers:
+                if hasattr(worker, "force_assign_with"):
+                    for shift_name, with_shifts in getattr(worker, "force_assign_with", {}).items():
+                        missing = set([shift_name] + list(with_shifts)) - all_shift_names
+                        if missing:
+                            raise InvalidShift(
+                                f"Worker '{worker.name}' has force_assign_with referencing non-existent shift(s): {missing}"
+                            )
+
             # Check worker requirements
             if not isinstance(s.workers_required, int):
                 # Validate the worker requirements
@@ -4166,22 +4234,36 @@ class RotaBuilder(object):
 
         return shifts
 
-    def get_worker_shift_list(self, worker: Worker, include_locums=False) -> List:
+    def get_worker_shift_list(self, worker: Worker, include_locums=False, search_multiple_assignments=False) -> List:
+        """
+        This function returns a list of shifts assigned to a worker for each week and day.
+
+        It was originally designed to return a list of shift names (when only a single shift per day was supported), but it can also return a set of shift names if `search_multiple_assignments` is True.
+        """
         shifts = []
 
         for week, day in self.get_week_day_combinations():
             # d = self.start_date + datetime.timedelta(n)
             # n = n + 1
 
-            temp = ""
+            if search_multiple_assignments:
+                temp = set()
+            else:
+                temp = ""
             for shift in self.get_shift_names_by_week_day(week, day):
                 if self.model.works[worker.id, week, day, shift].value > 0.90:
-                    temp = shift
+                    if search_multiple_assignments:
+                        temp.add(shift)
+                    else:
+                        temp = shift
                 elif (
                     include_locums
                     and self.model.locum_works[worker.id, week, day, shift].value > 0.90
                 ):
-                    temp = shift
+                    if search_multiple_assignments:
+                        temp.add(shift)
+                    else:
+                        temp = shift
             shifts.append(temp)
 
         return shifts

rota/workers.py

@@ -114,13 +114,17 @@ class Worker(BaseModel):
     allowed_multi_shift_sets: list[set] = []  # or list/set of frozensets
     prefer_multi_shift_together: int = 0  # Default: no preference
     # Set to a positive integer to prefer, negative to discourage, 0 for neutral
-    hard_day_dependencies: dict[str, str] = {}  # e.g. {"Fri": "Sun"}
+    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"]}
+
 
     assign_as_block_preferences: dict[str, float] = {}  # shift_name -> weight (positive = prefer block)
 
     shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist
 
+    weekend_shift_target_number: int = 0
+
     model_config = ConfigDict(
         extra="allow",
     )
@@ -402,12 +406,19 @@ class Worker(BaseModel):
             self.allowed_multi_shift_sets = []
         self.allowed_multi_shift_sets.append(frozenset(shifts))
 
-    def add_hard_day_dependency(self, from_day: str, to_day: str):
+    def add_hard_day_dependency(self, *days: str):
         if not hasattr(self, "hard_day_dependencies"):
-            self.hard_day_dependencies = {}
+            self.hard_day_dependencies = set()
-        self.hard_day_dependencies[from_day] = to_day
+        self.hard_day_dependencies.add(frozenset(days))
 
     def add_hard_day_exclusion(self, day1: str, day2: str):
         if not hasattr(self, "hard_day_exclusions"):
             self.hard_day_exclusions = []
-        self.hard_day_exclusions.append((day1, day2))
+        self.hard_day_exclusions.append((day1, day2))
+
+    def add_force_assign_with(self, shift: str, with_shifts: list[str] | str):
+        if not hasattr(self, "force_assign_with"):
+            self.force_assign_with = {}
+        if isinstance(with_shifts, str):
+            with_shifts = [with_shifts]
+        self.force_assign_with[shift] = with_shifts

test/test_workers.py

@@ -1,5 +1,5 @@
 import pytest
-from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days
 
 import datetime
 from rota.workers import NonWorkingDays, NotAvailableToWork, Worker, generate_not_available_to_works
@@ -1169,4 +1169,214 @@ def test_worker_hard_day_dependency_weekend2():
                 assert shift_list[sat_idx] == "a", f"Worker4 assigned 'a' on Fri (week {week}) but not on Sun"
 
 
+def test_worker_force_assign_with_invalid_shift():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=2)
+    workers = Rota.get_workers()
+    for worker in workers:
+        worker.add_force_assign_with("a", "c")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+
+    with pytest.raises(InvalidShift):
+        # This should raise an InvalidShift error because 'c' does not exist
+        Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="c",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
             
+
+def test_worker_force_assign_with():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_force_assign_with("a", "b")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign_with", folder="tests")
+
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker, search_multiple_assignments=True)
+
+        for shift in shift_list:
+            if shift == "a":
+                assert "b" in shift_list, f"Worker {worker.name} has 'a' but not 'b' assigned with it"
+            elif shift == "b":
+                assert "a" in shift_list, f"Worker {worker.name} has 'b' but not 'a' assigned with it"
+
+        assert shift_list.count({"a", "b"}) == 16
+
+def test_worker_force_assign_with_and_hard_days():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_force_assign_with("a", "b")
+
+        #worker.add_hard_day_dependency("Sat", "Sun")
+        worker.add_hard_day_dependency("Mon", "Tue")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:2],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:2],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign_with_hard_days", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+
+def test_worker_force_assign_with_and_hard_days_weekend():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_force_assign_with("a", "b")
+
+        #worker.add_hard_day_dependency("Sat", "Sun")
+        worker.add_hard_day_dependency("Sat", "Sun")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[5:],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[5:],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign_with_hard_days", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+
+
+def test_workers_double_shifts_on_mon_tue():
+    Rota = generate_basic_rota(workers=2, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_force_assign_with("a", "b")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:2],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:2],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+
+def test_workers_double_shifts_on_weekends():
+    Rota = generate_basic_rota(workers=2, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_force_assign_with("a", "b")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[5:],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[5:],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
+
+    assert Rota.results.solver.status == "ok"