...

output/timetable.css

@@ -318,3 +318,7 @@ table.transposed th.bank-holiday {
 .target-assigned.no-colour {
     background-color: white !important;
 }
+
+.locum-shift {
+    background-color: lightgreen;
+}

output/timetable.js

@@ -283,6 +283,7 @@ $(".table-div .worker-row .worker").each((n, tr) => {
     }
 
     shift_counts = jtr.data("shift-counts")
+    locum_shift_counts = jtr.data("locum-shift-counts")
     shift_targets = jtr.data("worker-targets")
 
     oshifts.forEach((s) => {
@@ -292,7 +293,11 @@ $(".table-div .worker-row .worker").each((n, tr) => {
             } else {
                 c = 0;
             }
-            row.append(`<td class="target-assigned ${s}" data-assigned=${c} data-target=${shift_targets[s].toFixed(2)} data-diff=${shift_targets[s].toFixed(2)-c}>${c} (${shift_targets[s].toFixed(2)})</td>`)
+            locum_shifts = ""
+            if (s in locum_shift_counts) {
+                locum_shifts = `+${locum_shift_counts[s]}`
+            }
+            row.append(`<td class="target-assigned ${s}" data-assigned=${c} data-target=${shift_targets[s].toFixed(2)} data-diff=${shift_targets[s].toFixed(2)-c}>${c}${locum_shifts} (${shift_targets[s].toFixed(2)})</td>`)
 
         } else {
             row.append(`<td>-</td>`)

pytest.ini

@@ -0,0 +1,4 @@
+# pytest.ini  
+[pytest]  
+testpaths = test  
+python_files = test_*.py

rota/shifts.py

@@ -268,9 +268,11 @@ class RotaBuilder(object):
             "balance_bank_holidays": True,
             "balance_blocks": True,
             "balance_shifts": True,  # Does not use a quadratic function
+            "balance_shifts_true_quadratic": False,  # Does not use a quadratic function
             "balance_shifts_quadratic": False,  # Will prevent spreading of spreading across different shifts
             "balance_shifts_over_workers": True,
             "minimise_shift_diffs": False,  # less sophisticated version of balance_shifts_over_workers
+            "maximum_allowed_shift_diff": None,
             "balance_weekends": True,
             "max_weekends": 100,
             # Don't assign multiple shifts every (n) weeks
@@ -289,6 +291,11 @@ class RotaBuilder(object):
             "hard_constrain_pair_separation": False,
             "avoid_shifts_by_grades": [],
             "avoid_shifts_by_worker_names": [],
+            "allocate_locum_shifts": True,
+            "distribute_locum_shifts": True,
+            "balance_locum_shifts": False,
+            "maximum_allowed_locum_shifts_per_worker": None,
+            "minimum_allowed_locum_shifts_per_worker": None
         }
 
         self.terminate_on_warning = [
@@ -297,6 +304,8 @@ class RotaBuilder(object):
             "Worker/no valid shifts",
             "Shift/invalid start date",
             "Shift/invalid end date",
+            "Locum/no locum availability",
+
         ]
 
         self.results = None
@@ -864,6 +873,8 @@ class RotaBuilder(object):
 
 
         def locum_request_init(model, wid, week, day, shift):
+            print("locum rquest init")
+            print(locals())
             if (wid, week, day, shift) in self.locum_availability:
                 self.locum_availability_map[(wid, week, day)] = shift
                 return 1
@@ -873,6 +884,7 @@ class RotaBuilder(object):
             (
                 (worker.id, week, day, shift)
                 for worker in self.workers
+                if not worker.locum
                 for week, day, shift in self.get_all_shiftname_combinations()
             ),
             initialize=locum_request_init,
@@ -931,6 +943,59 @@ class RotaBuilder(object):
         # binary variables representing if a worker worked on sunday but not on saturday (avoid if possible)
         # self.model.no_pref = Var([worker.id for worker in self.workers], within=Binary, initialize=0)
 
+        if self.get_workers_who_require_locums():
+            # We use this as a time check that their is some valid locum ability
+            if not self.get_all_locum_availability():
+                self.add_warning("Locum/no locum availability", "No locum availability set")
+
+
+
+            self.model.locum_required = Var(
+                (
+                    (week, day, shiftname)
+                    for week, day, shiftname in self.get_all_shiftname_combinations()
+                ),
+                within=Binary,
+                initialize=0,
+            )
+
+            self.model.locum_works = Var(
+                (
+                    (worker.id, week, day, shiftname)
+                    for worker in self.workers
+                    for week, day, shiftname in self.get_all_shiftname_combinations()
+                ),
+                within=Binary,
+                initialize=0,
+            )
+
+            self.model.locum_shifts_by_worker = Var(
+                (
+                    worker.id
+                    for worker in self.workers
+                ),
+                within=Integers,
+                initialize=0,
+            )
+
+            self.model.locum_shifts_t1 = Var(
+                (
+                    worker.id
+                    for worker in self.workers
+                ),
+                domain=NonNegativeReals,
+                initialize=0,
+            )
+
+            self.model.locum_shifts_t2 = Var(
+                (
+                    worker.id
+                    for worker in self.workers
+                ),
+                domain=NonNegativeReals,
+                initialize=0,
+            )
+
         self.build_model_constraints()
 
     def build_model_constraints(self):
@@ -953,6 +1018,36 @@ class RotaBuilder(object):
                     if worker.site in site_required
                 )
             )
+
+            if self.get_workers_who_require_locums():
+                self.model.constraints.add(
+                    self.model.locum_required[week, day, shift]
+                    == sum(
+                        self.model.works[worker.id, week, day, shift]
+                        for worker in self.workers
+                        if worker.locum
+                    )
+                )
+
+                self.model.constraints.add(
+                    self.model.locum_required[week, day, shift]
+                    == sum(
+                        self.model.locum_works[worker.id, week, day, shift]
+                        for worker in self.workers
+                        #if not worker.locum
+                    )
+                )
+                self.model.constraints.add(
+                    0
+                    == sum(
+                        self.model.locum_works[worker.id, week, day, shift]
+                        for worker in self.workers
+                        if worker.locum
+                    )
+                )
+
+
+
             # 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]:
                 self.model.constraints.add(
@@ -1190,20 +1285,6 @@ class RotaBuilder(object):
             for shift_name in self.shifts_to_assign_or_force_as_blocks():
                 shift = self.get_shift_by_name(shift_name)
 
-                print(shift_name)
-                print(
-                    week,
-                    self.get_week_start_date(week),
-                    shift.start_date,
-                    shift.end_date,
-                )
-
-                # if self.get_week_start_date(week) <= shift.start_date:
-                #    continue
-
-                # if self.get_week_start_date(week) > shift.end_date:
-                #    continue
-
                 for worker in self.get_workers_for_shift(shift):
                     try:
                         self.model.constraints.add(
@@ -1231,7 +1312,6 @@ class RotaBuilder(object):
                 workers_required = shift.get_worker_requirement_by_date(
                     self.get_week_start_date(week)
                 )
-                print(worker.name, week, workers_required, self.get_week_start_date(week))
 
                 if shift.force_as_block_unless_nwd:
                     workers = self.get_workers_for_shift(shift)
@@ -1304,7 +1384,70 @@ class RotaBuilder(object):
             )
             logging.debug(f"Generate worker constraints: {worker.name}")
 
-            if self.constraint_options["minimise_shift_diffs"]:
+
+
+            if self.get_workers_who_require_locums():
+                for week, day, shift in self.get_all_shiftclass_combinations():
+                    if not worker.locum:
+                        self.model.constraints.add(
+                            self.get_locum_availability(worker, week, day, shift)
+                            >= 
+                                self.model.locum_works[worker.id, week, day, shift.name]
+                            
+                        )
+
+                        self.model.constraints.add(
+                            self.model.locum_shifts_by_worker[worker.id] ==
+                            sum(
+                                self.model.locum_works[worker.id, week, day, shift]
+                                for week, day, shift in self.get_all_shiftname_combinations()
+                            )
+                        )
+
+
+                self.model.constraints.add(
+                    self.model.locum_shifts_t1[worker.id]
+                    - self.model.locum_shifts_t2[worker.id]
+                    == self.model.locum_shifts_by_worker[worker.id]
+                )
+
+                if worker.locum_max_shifts:
+                    self.model.constraints.add(
+                        worker.locum_max_shifts
+                        >= 
+                            sum(self.model.locum_works[worker.id, week, day, shiftname] for week, day, shiftname in self.get_all_shiftname_combinations())
+                        
+                    )
+
+                if worker.locum_max_shifts_per_week:
+                    for week_to_check in self.weeks:
+                        self.model.constraints.add(
+                            worker.locum_max_shifts_per_week
+                            >= 
+                                sum(self.model.locum_works[worker.id, week, day, shiftname] for week, day, shiftname in self.get_all_shiftname_combinations(week=week_to_check))
+                            
+                        )
+
+                if self.constraint_options["distribute_locum_shifts"]:
+                    self.model.constraints.add(
+                        sum(self.model.locum_required[week, day, shift] for week, day, shift in self.get_all_shiftname_combinations()) / (len(self.get_locum_workers())-1)
+                        >= self.model.locum_shifts_by_worker[worker.id]
+                    )
+
+                if self.constraint_options["maximum_allowed_locum_shifts_per_worker"] is not None:
+                    self.model.constraints.add(
+                        self.constraint_options["maximum_allowed_locum_shifts_per_worker"]
+                        >= self.model.locum_shifts_by_worker[worker.id]
+                    )
+                if worker.locum_availability:
+                    if self.constraint_options["minimum_allowed_locum_shifts_per_worker"] is not None:
+                        self.model.constraints.add(
+                            self.constraint_options["minimum_allowed_locum_shifts_per_worker"]
+                            <= self.model.locum_shifts_by_worker[worker.id]
+                        )
+
+
+            if self.constraint_options["minimise_shift_diffs"] or self.constraint_options["maximum_allowed_shift_diff"] is not None:
                 self.model.constraints.add(
                     self.model.shift_count_diff_summed[worker.id]
                     == sum(
@@ -1313,6 +1456,16 @@ class RotaBuilder(object):
                     )
                 )
 
+            if self.constraint_options["maximum_allowed_shift_diff"] is not None:
+                self.model.constraints.add(
+                    self.model.shift_count_diff_summed[worker.id]
+                    <= self.constraint_options["maximum_allowed_shift_diff"]
+                )
+                self.model.constraints.add(
+                    self.model.shift_count_diff_summed[worker.id]
+                    >= -self.constraint_options["maximum_allowed_shift_diff"]
+                )
+
             for week_blocks in self.get_week_block_iterator(4):
                 # Prevent more than n number shifts per 4 weeks
                 try:
@@ -1743,6 +1896,14 @@ class RotaBuilder(object):
                                         worker.id, week, day, shift.name
                                     ]
                                 )
+                        if not worker.locum_on_nwds and day == n:
+                            if start <= self.week_day_date_map[(week, day)] < end:
+                                self.model.constraints.add(
+                                    0
+                                    == self.model.locum_works[
+                                        worker.id, week, day, shift.name
+                                    ]
+                                )
 
             if self.constraint_options["balance_blocks"]:
                 if self.constraint_options["max_night_frequency"]:
@@ -2231,6 +2392,14 @@ class RotaBuilder(object):
                             for shift in self.get_shift_names_by_week_day(week, day)
                         )
                     )
+                    if self.get_locum_workers():
+                        self.model.constraints.add(
+                            self.model.available[worker.id, week, day]
+                            >= sum(
+                                self.model.locum_works[worker.id, week, day, shift]
+                                for shift in self.get_shift_names_by_week_day(week, day)
+                            )
+                        )
 
                     # single shift per day
                     self.model.constraints.add(
@@ -2240,6 +2409,22 @@ class RotaBuilder(object):
                             for shift in self.get_shift_names_by_week_day(week, day)
                         )
                     )
+                    # This applies to locums as well
+                    if self.get_workers_who_require_locums():
+                        self.model.constraints.add(
+                            1
+                            >= sum(
+                                self.model.locum_works[worker.id, week, day, shift]
+                                for shift in self.get_shift_names_by_week_day(week, day)
+                            )
+                        )
+                        self.model.constraints.add(
+                            1
+                            >= sum(
+                                self.model.works[worker.id, week, day, shift] + self.model.locum_works[worker.id, week, day, shift]
+                                for shift in self.get_shift_names_by_week_day(week, day)
+                            )
+                        )
 
                     if self.constraint_options["hard_constrain_pair_separation"]:
                         for worker_pairs in self.worker_pairs:
@@ -2349,78 +2534,6 @@ class RotaBuilder(object):
                             # except KeyError:
                             #    pass
 
-                #                for constraint_shift in self.get_shifts_with_constraints(
-                #                    "preclear", "preclear2"
-                #                ):
-                #                    if day in constraint_shift.days:
-                #                        self.model.constraints.add(
-                #                            1
-                #                            >= self.model.works[
-                #                                worker.id, week, day, constraint_shift.name
-                #                            ]
-                #                            + sum(
-                #                                p1 * self.model.works[w.id, pweek, pday, shiftname]
-                #                                for shiftname in self.get_shift_names_by_week_day(
-                #                                    pweek, pday
-                #                                )
-                #                                if shiftname != constraint_shift.name
-                #                                for w in workers
-                #                            )
-                #                        )
-                #
-                #                for constraint_shift in self.get_shifts_with_constraint("preclear2"):
-                #                    if day in constraint_shift.days:
-                #                        self.model.constraints.add(
-                #                            1
-                #                            >= self.model.works[
-                #                                worker.id, week, day, constraint_shift.name
-                #                            ]
-                #                            + sum(
-                #                                p2 * self.model.works[w.id, p2week, p2day, shiftname]
-                #                                for shiftname in self.get_shift_names_by_week_day(
-                #                                    p2week, p2day
-                #                                )
-                #                                if shiftname != constraint_shift.name
-                #                                for w in workers
-                #                            )
-                #                        )
-
-                # for constraint_shift in self.get_shifts_with_constraints(
-                #    "postclear", "postclear2"
-                # ):
-                #    if day in constraint_shift.days:
-                #        self.model.constraints.add(
-                #            1
-                #            >= self.model.works[
-                #                worker.id, week, day, constraint_shift.name
-                #            ]
-                #            + sum(
-                #                n1 * self.model.works[w.id, nweek, nday, shiftname]
-                #                for shiftname in self.get_shift_names_by_week_day(
-                #                    nweek, nday
-                #                )
-                #                if shiftname != constraint_shift.name
-                #                for w in workers
-                #            )
-                #        )
-
-                # for constraint_shift in self.get_shifts_with_constraints("postclear2"):
-                #    if day in constraint_shift.days:
-                #        self.model.constraints.add(
-                #            1
-                #            >= self.model.works[
-                #                worker.id, week, day, constraint_shift.name
-                #            ]
-                #            + sum(
-                #                n2 * self.model.works[w.id, n2week, n2day, shiftname]
-                #                for shiftname in self.get_shift_names_by_week_day(
-                #                    n2week, n2day
-                #                )
-                #                if shiftname != constraint_shift.name
-                #                for w in workers
-                #            )
-                #        )
-
                 # Night constraint means we won't assign a shift the day before
                 # an unavailability
                 for constraint_shift in self.get_shifts_with_constraint("night"):
@@ -2444,6 +2557,32 @@ class RotaBuilder(object):
                             ]
                         )
 
+        ## Ensure each shift has the requisit number of workers assigned
+        #for (
+        #    week,
+        #    day,
+        #    shift,
+        #    workers_required,
+        #    site_required,
+        #) in self.get_required_workers_and_site_combinations():
+        #    # print(week, day, shift, workers_required, site_required)
+        #    self.model.constraints.add(
+        #        workers_required
+        #        == sum(
+        #            self.model.works[worker.id, week, day, shift]
+        #            for worker in self.workers
+        #            if worker.locum
+        #        )
+        #    )
+
+        #for worker in self.workers:
+        #    if worker.locum:
+        #        for week, day, shift in self.get_all_shiftname_combinations():
+        #            self.model.constraints.add(
+        #                self.model.locum_works[worker.id, week, day, shift]
+        #                == self.model.works[worker.id, week, day, shift]
+        #            )
+
         self.define_objectives()
 
         print("Building model completed")
@@ -2453,9 +2592,20 @@ class RotaBuilder(object):
         def obj_rule(m):
             # c = len(workers)
 
-            balance_modifier_constant = 1
+            balance_modifier_constant = 10
             balance_quadratic_shift_modifier_constant = 3
             block_shift_balancing_constant = 1
+            locum_shift_balance_modifier_constant = 40
+
+            locum_shift_balancing = 0
+            if self.get_workers_who_require_locums():
+                if self.constraint_options["balance_locum_shifts"]:
+                    locum_shift_balancing = sum(
+                        locum_shift_balance_modifier_constant * 
+                        (self.model.locum_shifts_t1[(worker.id)]
+                        + self.model.locum_shifts_t2[(worker.id)])
+                        for worker in self.workers
+                    )
 
             if self.constraint_options["balance_shifts_over_workers"]:
                 worker_shift_balancing = sum(
@@ -2495,6 +2645,16 @@ class RotaBuilder(object):
             else:
                 quadratic_shift_balancing = 0
 
+            true_quadratic_shift_balancing = 0
+            if self.constraint_options["balance_shifts_true_quadratic"]:
+                shift_diff_modifier_constant = 10
+                true_quadratic_shift_balancing = sum(
+                    shift_diff_modifier_constant
+                    * self.model.shift_count_diff_summed[(worker.id)]
+                    for worker in self.workers
+                )
+            
+
             if self.constraint_options["minimise_shift_diffs"]:
                 shift_diff_modifier_constant = 10
                 shift_diff_balancing = sum(
@@ -2622,6 +2782,8 @@ class RotaBuilder(object):
                 + bank_holiday_balancing
                 + blocks_balancing
                 - work_requests
+                + locum_shift_balancing
+                + true_quadratic_shift_balancing
             )
 
         # add objective function to the model. rule (pass function) or expr (pass expression directly)
@@ -2796,6 +2958,18 @@ class RotaBuilder(object):
             raise ValueError(f"Must pair at least two shifts: {shifts}")
         self.paired_shifts.append([*shifts])
 
+    def get_date_range(self, start_date: datetime.date=None, end_date: datetime.date = None):
+        """Gets a range of dates
+        
+        If either start_date or end_date are not provided defaults to the rota dates"""
+        if start_date is None:
+            start_date = self.start_date
+        if end_date is None:
+            end_date = self.rota_end_date
+
+        for n in range(int((end_date - start_date).days) + 1):
+            yield start_date + datetime.timedelta(n)
+
     def build_shifts(self):
         """
         Process the added shifts
@@ -3145,6 +3319,16 @@ class RotaBuilder(object):
         s.extend(self.shifts_to_force_as_blocks())
         return s
     
+    def get_all_locum_availability(self):
+        return self.locum_availability_map
+
+    def get_locum_availability(self, worker: Worker, week: int, day: DayStr, shift):
+        
+        if (worker.id, week, day) in self.locum_availability_map:
+            if shift.name in self.locum_availability_map[(worker.id, week, day)]:
+                return 1
+        return 0
+
     def get_all_workers(self) -> List[Worker]:
         return self.workers
 
@@ -3200,6 +3384,12 @@ class RotaBuilder(object):
     def get_worker_grades(self) -> set[int]:
         return set([worker.grade for worker in self.workers])
 
+    def get_workers_who_require_locums(self):
+        return [worker for worker in self.get_all_workers() if worker.locum]
+
+    def get_locum_workers(self):
+        return [worker for worker in self.get_all_workers() if not worker.locum and worker.locum_availability]
+
     def get_bank_holiday_week_days(self):
         return [
             (week, day)
@@ -3400,6 +3590,7 @@ class RotaBuilder(object):
                 current_site = worker.site
 
             shifts = []
+            locum_shifts = []
             nwds = json.dumps(worker.non_working_day_list, default=str)
             # if worker.nwd:
             #    # TODO: limit to dates
@@ -3423,12 +3614,22 @@ class RotaBuilder(object):
 
                 # Loop through all the days possible shifts and see
                 # if the worker has been assigned
+                locum = False
                 for shift in self.get_shift_names_by_week_day(week, day):
-                    if model.works[worker.id, week, day, shift].value > 0.5:
+                    if model.works[worker.id, week, day, shift].value > 0.8:
                         shifts.append(shift)
                         a = shift[0]
                         shift_name = shift
                         break
+                    if self.get_workers_who_require_locums():
+                        if model.locum_works[worker.id, week, day, shift].value > 0.8:
+                            locum_shifts.append(shift)
+                            a = shift[0]
+                            shift_name = shift
+                            locum = True
+
+                            break
+                    
                 title = f"{shift_name} ({d})"
                 css_class = day
                 unavailable_reason = ""
@@ -3460,6 +3661,14 @@ class RotaBuilder(object):
                     title = " ".join((title, "[REQUESTED]"))
                     requests = f" data-shift-request='{self.work_requests_map[worker.id, week, day]}'"
 
+                if (worker.id, week, day) in self.locum_availability_map:
+                    css_class = " ".join((css_class, "locum-availability"))
+                    title = " ".join((title, "[LOCUM AVAILABILITY]"))
+                    requests = f" data-locum-request='{self.locum_availability_map[worker.id, week, day]}'"
+
+                if locum:
+                    css_class = " ".join((css_class, "locum-shift"))
+
                 remote_site = ""
                 if shift_name:
                     shift = self.get_shift_by_name(shift_name)
@@ -3479,6 +3688,13 @@ class RotaBuilder(object):
                 shift_count_dict[s] = c
                 shift_count = shift_count + f"{s}: {c}, "
 
+            locum_shift_count = ""
+            locum_shift_count_dict = {}
+            for s in set(locum_shifts):
+                c = locum_shifts.count(s)
+                locum_shift_count_dict[s] = c
+                locum_shift_count = locum_shift_count + f"{s}: {c}, "
+
             shift_diff_dict = {}
             for shift in self.get_shifts():
                 diff = model.shift_count_diff[worker.id, shift.name].value
@@ -3491,7 +3707,9 @@ class RotaBuilder(object):
             data-start_date='{start_date}' 
             data-end_date='{end_date}' 
             data-oops='{oops}' 
-            data-worker-targets='{targets}' data-shift-counts='{worker_shift_counts}' 
+            data-worker-targets='{targets}' 
+            data-shift-counts='{worker_shift_counts}' 
+            data-locum-shift-counts='{locum_shift_counts}'
             data-weekend-target='{weekend_target}' 
             data-remote-site='{remote_site}' 
             data-pair='{pair}'
@@ -3500,7 +3718,7 @@ class RotaBuilder(object):
             data-shift-diff='{shift_diff}'
             data-shift-diff-summed='{shift_diff_summed}'
             >
-            <span class='name' title='{name}'>{name}</span> ({grade}) [{fte}]</td>""".format(
+            <span class='name' title='{name} ({worker_id})'>{name}</span> ({grade}) [{fte}]</td>""".format(
                 site=worker.site,
                 nwds=nwds,
                 worker_id=worker.id,
@@ -3512,6 +3730,7 @@ class RotaBuilder(object):
                 oops=json.dumps(json.dumps(worker.oop, default=str)),
                 targets=worker_targets,
                 worker_shift_counts=json.dumps(shift_count_dict),
+                locum_shift_counts=json.dumps(locum_shift_count_dict),
                 weekend_target=worker.weekend_shift_target_number,
                 remote_site=worker.remote_site,
                 grade=worker.grade,
@@ -3617,6 +3836,17 @@ class RotaBuilder(object):
         <div id="shift-timetable-div">
         <div>
         </details>
+        <details>
+        <summary><h2>Extra</h2></summary>
+        <div id="extra-options">
+        </div>
+        <div id="extra-div">
+        <h3>Locum availability</h3>
+        {"TEST"}
+        {self.locum_availability_map}
+        <br/>
+        <div>
+        </details>
         </div
         </body>
         """
@@ -3674,7 +3904,7 @@ class RotaBuilder(object):
 
         return shifts
 
-    def get_worker_shift_list(self, worker: Worker) -> List:
+    def get_worker_shift_list(self, worker: Worker, include_locums=False) -> List:
         shifts = []
 
         for week, day in self.get_week_day_combinations():
@@ -3685,6 +3915,8 @@ class RotaBuilder(object):
             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 include_locums and self.model.locum_works[worker.id, week, day, shift].value > 0.90:
+                    temp = shift
             shifts.append(temp)
 
         return shifts

rota/workers.py

@@ -67,6 +67,9 @@ class NonWorkingDays(BaseModel):
 
 
 class WorkRequests(BaseModel):
+    """ 
+    Shift is the name of the shift requested, * is a wildcard for any shift
+    """
     date: datetime.date
     shift: str = ""
 
@@ -99,6 +102,9 @@ class Worker(BaseModel):
     pref_not_to_work: list[PreferenceNotToWork] = []
     work_requests: list[WorkRequests] = []
     locum_availability: list[WorkRequests] = []
+    locum_max_shifts: int = 0
+    locum_max_shifts_per_week: int = 0
+    locum_on_nwds: bool = False
     remote_site: str = "plymouth"  # We set a default proc_site
     previous_shifts: dict = {}
     shift_balance_extra: dict = {}
@@ -293,7 +299,8 @@ class Worker(BaseModel):
             week = days_from_start // 7 + 1
             day = Rota.days[(days_from_start % 7)]
             if request.shift == "*":
-                for shift in Rota.get_shifts_for_worker_site(self.site):
+                # We may want to add a limit to shifts that can be filled by locums?
+                for shift in Rota.get_shifts():
                     Rota.locum_availability.add((self.id, week, day, shift.name))
             else:
                 Rota.locum_availability.add((self.id, week, day, request.shift))

test/test_shifts.py

@@ -591,7 +591,7 @@ class TestShiftWorkerRequirements:
             SingleShift(
                 sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
                 assign_as_block=True,
-                force_as_block=True,
+                #force_as_block=True,
                 workers_required=wr,
                 #start_date=(Rota.start_date + datetime.timedelta(days=7)),
                 #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
@@ -619,7 +619,7 @@ class TestShiftWorkerRequirements:
             SingleShift(
                 sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
                 assign_as_block=True,
-                force_as_block=True,
+                #force_as_block=True,
                 workers_required=wr,
                 #start_date=(Rota.start_date + datetime.timedelta(days=7)),
                 #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
@@ -686,7 +686,7 @@ class TestShiftWorkerRequirements:
             SingleShift(
                 sites=("group1", "group2"), name="b", length= 12.5, days=days[4:],
                 assign_as_block=True,
-                force_as_block=True,
+                #force_as_block=True,
                 workers_required=2,
                 #start_date=(Rota.start_date + datetime.timedelta(days=7)),
                 #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
@@ -694,7 +694,7 @@ class TestShiftWorkerRequirements:
             SingleShift(
                 sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
                 assign_as_block=True,
-                force_as_block=True,
+                #force_as_block=True,
                 workers_required=wr,
                 #start_date=(Rota.start_date + datetime.timedelta(days=7)),
                 #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),

test/test_worker_requests.py

@@ -95,7 +95,7 @@ class TestWorkerRequests:
         )
 
         Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("preferences_not_to_work2")
+        Rota.export_rota_to_html("preferences_not_to_work2", timestamp_filename=False)
 
         shift_summary = Rota.get_shift_summary_dict()
         for worker_name in shift_summary:

test/test_workers.py

@@ -2,7 +2,8 @@ import pytest
 from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days
 
 import datetime
-from rota.workers import NotAvailableToWork, Worker, generate_not_available_to_works
+from rota.workers import NonWorkingDays, NotAvailableToWork, Worker, generate_not_available_to_works
+from rota.workers import WorkRequests
 
 
 def generate_basic_rota(
@@ -619,6 +620,7 @@ class TestComplexWorkers:
         #            assert worker.get_fte("b") == 50
         #            assert total_shifts[worker.name] == 20
         #total_shifts = Rota.get_workers_total_shifts()
+        assert Rota.results.solver.status == "ok"
         for worker in Rota.get_workers():
             shift_string = Rota.get_worker_shift_list_string(worker)
             assert shift_string.count("cccc") == shift_string.count("c") / 4
@@ -674,30 +676,198 @@ class TestComplexWorkers:
         Rota.build_and_solve()
 
         Rota.export_rota_to_html("test_split")
+        assert Rota.results.solver.status == "ok"
 
 
 class TestLocums:
-    def test_locums(self):
+    def test_locums_basic(self):
-        Rota = generate_basic_rota(workers=9, weeks_to_rota=22)
+        Rota = generate_basic_rota(workers=2, weeks_to_rota=6)
 
-        Rota.set_rota_constraint("max_shifts_per_week", 2)
+        la = []
-        Rota.set_rota_constraint("max_weekend_frequency", 3)
+        for date in Rota.get_date_range(end_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
-        Rota.set_rota_constraint("max_shifts_per_month", 4)
+        #for date in Rota.get_date_range():
-
+            la.append(WorkRequests(date=date, shift="*"))
-        Rota.terminate_on_warning.remove("Worker/no valid shifts")
+        la2 = []
+        for date in Rota.get_date_range(start_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
+        #for date in Rota.get_date_range():
+            la2.append(WorkRequests(date=date, shift="*"))
 
         Rota.add_workers(
             [
+                Worker(name="worker03", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1),
+                Worker(name="worker04", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2),
                 Worker(name="locum1", site="group1", locum=True),
                 Worker(
                     name="locum2", site="group1", locum=True
                 ),
-                Worker(
+            ]
-                    name="extra1", site="locum"
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:4],
+                workers_required=1,
+                
+                constraint=[
+                #{"name": "pre", "options": 1},
+                #{"name": "post", "options": 1},
+                #{"name": "max_shifts_per_week", "options": 1},
+                ]
             ),
+        )
+        Rota.set_rota_constraint("distribute_locum_shifts", True)
+        Rota.set_rota_constraint("maximum_allowed_shift_diff", 0)
+        Rota.set_rota_constraint("maximum_allowed_locum_shifts_per_worker", 4)
+
+        Rota.build_and_solve()
+
+        Rota.export_rota_to_html("test_locums_basic", timestamp_filename=False)
+
+        assert Rota.results.solver.status == "ok"
+
+        for worker in Rota.get_workers():
+            assert Rota.get_worker_shift_list(worker).count("a") == 4
+
+    def test_locums_nwds(self):
+        Rota = generate_basic_rota(workers=2, weeks_to_rota=6)
+
+        la = []
+        for date in Rota.get_date_range(end_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
+            la.append(WorkRequests(date=date, shift="*"))
+        la2 = []
+        for date in Rota.get_date_range(start_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
+            la2.append(WorkRequests(date=date, shift="*"))
+
+        Rota.add_workers(
+            [
+                Worker(name="worker03", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1, nwds=[NonWorkingDays(day="Tue")]),
+                Worker(name="worker04", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2, nwds=[NonWorkingDays(day="Wed"),NonWorkingDays(day="Thu")]),
+                #Worker(name="worker03", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1, nwds=[NonWorkingDays(day="Mon"),NonWorkingDays(day="Tue")]),
+                #Worker(name="worker04", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2, nwds=[NonWorkingDays(day="Wed"),NonWorkingDays(day="Thu")]),
+                Worker(name="locum1", site="group1", locum=True),
                 Worker(
-                    name="extra2", site="locum"
+                    name="locum2", site="group1", locum=True
                 ),
+            ]
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:4],
+                workers_required=1,
+                
+                constraint=[
+                #{"name": "pre", "options": 1},
+                #{"name": "post", "options": 1},
+                #{"name": "max_shifts_per_week", "options": 1},
+                ]
+            ),
+        )
+        #Rota.set_rota_constraint("distribute_locum_shifts", True)
+        Rota.set_rota_constraint("maximum_allowed_shift_diff", 0)
+        Rota.set_rota_constraint("maximum_allowed_locum_shifts_per_worker", 4)
+        #Rota.set_rota_constraint("minimum_allowed_locum_shifts_per_worker", 3)
+
+        Rota.build_and_solve()
+
+        Rota.export_rota_to_html("test_locums_basic", timestamp_filename=False)
+
+        assert Rota.results.solver.status == "ok"
+
+        for worker in Rota.get_workers():
+            print(  Rota.get_worker_shift_list(worker))
+            assert Rota.get_worker_shift_list(worker).count("a") == 4
+
+            if worker.name in ("worker03", "worker04"):
+                assert Rota.get_worker_shift_list(worker, include_locums=True).count("a") > 6
+
+    def test_locums_no_availablity(self):
+        Rota = generate_basic_rota(workers=2, weeks_to_rota=5)
+
+        la = []
+        #for date in Rota.get_date_range(end_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
+        #    la.append(WorkRequests(date=date, shift="*"))
+        la2 = []
+        #for date in Rota.get_date_range(start_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
+        #    la2.append(WorkRequests(date=date, shift="*"))
+
+        Rota.add_workers(
+            [
+                Worker(name="worker03", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1),
+                Worker(name="worker04", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2),
+                Worker(name="locum1", site="group1", locum=True),
+                Worker(
+                    name="locum2", site="group1", locum=True
+                ),
+            ]
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:4],
+                workers_required=1,
+                
+                constraint=[
+                #{"name": "pre", "options": 1},
+                #{"name": "post", "options": 1},
+                #{"name": "max_shifts_per_week", "options": 1},
+                ]
+            ),
+        )
+        Rota.set_rota_constraint("distribute_locum_shifts", False)
+
+        with pytest.raises(WarningTermination):
+            Rota.build_and_solve()
+
+
+
+    def test_locums(self):
+        Rota = generate_basic_rota(workers=7, weeks_to_rota=10)
+
+        Rota.set_rota_constraint("max_shifts_per_week", 2)
+        Rota.set_rota_constraint("max_weekend_frequency", 3)
+        Rota.set_rota_constraint("max_shifts_per_month", 4)
+        #Rota.set_rota_constraint("balance_shifts", False)
+        Rota.set_rota_constraint("distribute_locum_shifts", True)
+        Rota.set_rota_constraint("maximum_allowed_shift_diff", 1)
+        Rota.set_rota_constraint("maximum_allowed_locum_shifts_per_worker", 6)
+
+        Rota.terminate_on_warning.remove("Worker/no valid shifts")
+
+
+        la = []
+        for date in Rota.get_date_range(end_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
+            la.append(WorkRequests(date=date, shift="*"))
+        la2 = []
+        for date in Rota.get_date_range(start_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
+            la2.append(WorkRequests(date=date, shift="*"))
+
+
+        Rota.add_workers(
+            [
+                Worker(name="worker08", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1),
+                Worker(name="worker09", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2),
+                Worker(name="locum1", site="group1", locum=True),
+                Worker(
+                    name="locum2", site="group1", locum=True
+                ),
+                #Worker(
+                #    name="extra1", site="locum", locum_availability=la, 
+                #    locum_max_shifts=40
+                #),
+                #Worker(
+                #    name="extra2", site="locum", locum_availability=la,
+                #    locum_max_shifts=4, locum_max_shifts_per_week=1
+                #),
 
             ]
         )
@@ -710,6 +880,7 @@ class TestLocums:
                 length=12.5,
                 days=days[:4],
                 workers_required=1,
+                
                 constraint=[
                 {"name": "pre", "options": 2},
                 {"name": "post", "options": 3},
@@ -744,3 +915,5 @@ class TestLocums:
         Rota.build_and_solve()
 
         Rota.export_rota_to_html("test_locums", timestamp_filename=False)
+
+        assert Rota.results.solver.status == "ok"