From e1aec5011c4bd2810f9a101016dfe3ccf7cbe95a Mon Sep 17 00:00:00 2001
From: Ross <xkjq@ymail.com>
Date: Wed, 17 Sep 2025 21:35:52 +0100
Subject: [PATCH] Add option to allow force assignment with leave conflict in
 RotaBuilder

---
 gen_cons.py    |  1 +
 rota/shifts.py | 38 +++++++++++++++++++++++++++-----------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/gen_cons.py b/gen_cons.py
index de13146..d0a3506 100644
--- a/gen_cons.py
+++ b/gen_cons.py
@@ -322,6 +322,7 @@ def main(
         weeks_to_rota=weeks,
         balance_offset_modifier=bom,
         name="cons rota test",
+        allow_force_assignment_with_leave_conflict=True,
     )
     # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
     Rota.constraint_options["balance_weekends"] = True
diff --git a/rota/shifts.py b/rota/shifts.py
index 0af3e08..c9d4364 100644
--- a/rota/shifts.py
+++ b/rota/shifts.py
@@ -329,6 +329,7 @@ class RotaBuilder(object):
         use_previous_shifts: bool = False,
         use_shift_balance_extra: bool = False,
         use_bank_holiday_extra: bool = False,
+        allow_force_assignment_with_leave_conflict: bool = False,
         SHIFT_BOUNDS=SHIFT_BOUNDS,
         name: str = "",
         bank_holidays: dict[datetime.date, str] = bank_holiday_map,
@@ -422,12 +423,16 @@ class RotaBuilder(object):
         self.run_start_time: datetime.datetime | None = None
         self.run_end_time: datetime.datetime | None = None
 
+        self.allow_force_assignment_with_leave_conflict = allow_force_assignment_with_leave_conflict
+
         self.bank_holidays = bank_holidays
 
         self.paired_shifts = []
 
         self.exported_rota_file = None
 
+        self.MAX_SHIFTS_PER_DAY = 2
+
     def set_rota_constraint(self, constraint: str, value: Any):
         if constraint not in self.constraint_options:
             raise ValueError(f"Constraint {constraint} not valid")
@@ -962,12 +967,32 @@ class RotaBuilder(object):
         # self.model.unavailable = Var(((worker, week, day) for worker in self.workers for week in weeks for day in days),
         #                  within=Binary, initialize=0)
 
+
+        for worker in self.workers:
+            for week, day, shift_name in getattr(worker, "forced_assignments", []):
+                # Only add if this shift is valid for this week/day
+                if shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True):
+                    # Check for leave conflict
+                    if hasattr(self, "unavailable_to_work"):
+                        if (worker.id, week, day) in self.unavailable_to_work:
+                            if self.allow_force_assignment_with_leave_conflict:
+                                self.unavailable_to_work.remove((worker.id, week, day))
+                                self.add_warning(
+                                    "Force assignment/leave conflict",
+                                    f"Worker '{worker.name}' has a forced assignment for shift '{shift_name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which conflicts with a leave request. The leave request has been overridden.",
+                                )
+                            else:
+                                self.add_warning(
+                                    "Force assignment/leave conflict",
+                                    f"Worker '{worker.name}' has a forced assignment for shift '{shift_name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which conflicts with a leave request.",
+                                )
+
         def availability_init(model, wid, week, day):
             # This needs to be higher than the max number of possible shifts per day
             if (wid, week, day) in self.unavailable_to_work:
                 # print((wid, week, day))
                 return 0
-            return 3
+            return self.MAX_SHIFTS_PER_DAY + 1
 
         self.model.available = Param(
             (
@@ -1556,16 +1581,7 @@ class RotaBuilder(object):
             for week, day, shift_name in getattr(worker, "forced_assignments", []):
                 # Only add if this shift is valid for this week/day
                 if shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True):
-                    # Check for leave conflict
-                    leave_conflict = False
-                    if hasattr(self, "unavailable_to_work"):
-                        if (worker.id, week, day) in self.unavailable_to_work:
-                            leave_conflict = True
-                    if leave_conflict:
-                        self.add_warning(
-                            "Force assignment/leave conflict",
-                            f"Worker '{worker.name}' has a forced assignment for shift '{shift_name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which conflicts with a leave request.",
-                        )
+                    # We have already checked for leave conflicts when building the model
                     self.model.constraints.add(
                         self.model.works[worker.id, week, day, shift_name] == 1
                     )