e

rota/shifts.py

@@ -186,29 +186,10 @@ class RotaBuilder(object):
         self.SHIFT_BOUNDS = SHIFT_BOUNDS
         self.shifts: List[SingleShift] = []  # type List[SingleShift]
 
-        self.days = days
-
-        self.weeks = [i for i in range(1, weeks_to_rota + 1)]
-
-        self.weeks_days_product = list(itertools.product(self.weeks, days))
-
         self.use_previous_shifts = use_previous_shifts
         self.use_shift_balance_extra = use_shift_balance_extra
         self.use_bank_holiday_extra = use_bank_holiday_extra
 
-        if start_date.weekday() != 0:
-            raise ValueError(
-                f"Start date {start_date.isoformat()} must be a Mon (not a {days[start_date.weekday()]})"
-            )
-
-        self.start_date = start_date
-        self.rota_days_length = len(self.weeks) * 7
-        self.rota_end_date = self.start_date + datetime.timedelta(self.rota_days_length)
-
-        console.print(
-            f"Weeks to rota: {weeks_to_rota} ({self.start_date} - {self.rota_end_date})"
-        )
-
         self.unavailable_to_work = set()
         self.unavailable_to_work_reason = {}
         self.pref_not_to_work = {}
@@ -217,7 +198,6 @@ class RotaBuilder(object):
         self.work_requests_map = {}
 
         self.workers: list[Worker] = []
-        self.worker_pairs = []
 
         # self.night_blocks = ["weekday", "weekend", "none"]
 
@@ -242,6 +222,7 @@ class RotaBuilder(object):
             "max_weekend_frequency": 1,  # Requires balance weekends
             # Don't assign multiple shifts every (n) weeks
             "max_night_frequency": 1,
+            "max_night_frequency_week_exclusions": [],
             "max_shifts_per_week": 7,
             "max_shifts_per_month": 40,
             # The following will cause an unsolvable problem if a shift
@@ -255,6 +236,34 @@ class RotaBuilder(object):
             "avoid_shifts_by_grades": [],
         }
 
+
+        self.results = None
+
+        self.ignore_valid_shifts = False
+
+        self.set_rota_dates(start_date, weeks_to_rota)
+
+    def set_rota_dates(self, start_date, weeks_to_rota):
+
+        self.days = days
+
+        self.weeks = [i for i in range(1, weeks_to_rota + 1)]
+
+        self.weeks_days_product = list(itertools.product(self.weeks, days))
+
+        if start_date.weekday() != 0:
+            raise ValueError(
+                f"Start date {start_date.isoformat()} must be a Mon (not a {days[start_date.weekday()]})"
+            )
+
+        self.start_date = start_date
+        self.rota_days_length = len(self.weeks) * 7
+        self.rota_end_date = self.start_date + datetime.timedelta(self.rota_days_length)
+
+        console.print(
+            f"Weeks to rota: {weeks_to_rota} ({self.start_date} - {self.rota_end_date})"
+        )
+
         # Generate a map for week, day combinations to a given date
         self.week_day_date_map = {}
 
@@ -264,9 +273,6 @@ class RotaBuilder(object):
             self.week_day_date_map[(week, day)] = d
             n = n + 1
 
-        self.results = None
-
-        self.ignore_valid_shifts = False
 
     def solve_model(
         self,
@@ -307,6 +313,7 @@ class RotaBuilder(object):
         if not results.solver.status:
             sys.exit(0)
 
+
     def solve_shifts_individually(self, options, folder="individual_shifts"):
         shifts = self.shifts
 
@@ -1500,6 +1507,11 @@ class RotaBuilder(object):
                     for week_blocks in self.get_week_block_iterator(
                         self.constraint_options["max_night_frequency"]
                     ):
+                        # Ignore excluded weeks
+                        if set(week_blocks).intersection(set(self.constraint_options["max_night_frequency_week_exclusions"])):
+                            continue
+
+
                         if self.get_shifts_with_constraint("night"):
                             # Prevent nights more than once every n weeks
                             self.model.constraints.add(
@@ -2388,9 +2400,10 @@ class RotaBuilder(object):
                 w.fte_adj for w in self.workers if w.site == site
             )
 
+        self.worker_pairs = []
         pairs = defaultdict(list)
         for w in self.workers:
-            if w.pair > 0:
+            if w.pair is not None:
                 pairs[w.pair].append(w)
 
         for p in pairs:
@@ -2552,7 +2565,7 @@ class RotaBuilder(object):
 
         Returns:
             list: contains tuple of all possible week / day / shift combinations
-        """
+        """ 
         week_day_shifts = self.week_day_shift_product
 
         if week is not None:

rota/workers.py

@@ -81,7 +81,7 @@ class Worker(BaseModel):
     previous_shifts: dict = {}
     shift_balance_extra: dict = {}
     bank_holiday_extra: int = 0
-    pair: int = 0
+    pair: str | None = None
 
     class Config:
         extra = Extra.allow