diff --git a/rota/shifts.py b/rota/shifts.py index 1252b88..7b29cf1 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -252,6 +252,7 @@ class RotaBuilder(object): "avoid_st2_first_month": False, "hard_constrain_pair_separation": False, "avoid_shifts_by_grades": [], + "avoid_shifts_by_worker_names": [], } self.results = None @@ -1175,6 +1176,7 @@ class RotaBuilder(object): ) # Most of our constraints apply per worker + worker: Worker for worker in self.workers: console.rule( f"Generate worker constraints: [bold blue]{worker.name}[/bold blue]" @@ -1209,6 +1211,28 @@ class RotaBuilder(object): ) pass + for constraint_dict in self.constraint_options["avoid_shifts_by_worker_names"]: + if invalid_shifts := set(constraint_dict["shifts"]).difference( + set(self.get_shift_names()) + ): + if self.ignore_valid_shifts: + continue # Skip if we don't want to raise an error + else: + raise InvalidShift( + f"avoid shifts by grade constraint contains a non existent shift: {invalid_shifts}" + ) + + if worker.name in constraint_dict["names"]: + self.model.constraints.add( + 0 + == sum( + self.model.works[worker.id, week, day, shift] + for week, day, shift in self.get_all_shiftname_combinations() + if week in constraint_dict["weeks"] + and shift in constraint_dict["shifts"] + ) + ) + for constraint_dict in self.constraint_options["avoid_shifts_by_grades"]: if invalid_shifts := set(constraint_dict["shifts"]).difference( set(self.get_shift_names()) @@ -2522,6 +2546,13 @@ class RotaBuilder(object): for p in pairs: self.worker_pairs.append(tuple(pairs[p])) + def add_worker_name_constraint_by_week( + self, names: Iterable[str], weeks: Iterable[int], shifts + ): + self.constraint_options["avoid_shifts_by_worker_names"].append( + {"names": names, "weeks": weeks, "shifts": shifts} + ) + def add_grade_constraint_by_week( self, grades: Iterable[int], weeks: Iterable[int], shifts ): diff --git a/test/rota_constraints/test_avoid_workers_by_week.py b/test/rota_constraints/test_avoid_workers_by_week.py new file mode 100644 index 0000000..b857b2e --- /dev/null +++ b/test/rota_constraints/test_avoid_workers_by_week.py @@ -0,0 +1,196 @@ +import datetime +import pytest +from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days + +from rota.workers import Worker + +import itertools + + +def generate_basic_rota(weeks_to_rota=9): + start_date = datetime.date(2022, 3, 7) + + Rota = RotaBuilder( + start_date, + weeks_to_rota=weeks_to_rota, + ) + + # Add a few workers + Rota.add_workers( + [ + Worker(name="worker1", site="group1", grade=1), + Worker(name="worker2", site="group1", grade=2), + Worker(name="worker3", site="group1", grade=3), + ] + ) + + return Rota + + +def date_generator(from_date, days): + n = 0 + while True: + yield from_date + + n = n + 1 + if n >= days: + break + from_date = from_date + datetime.timedelta(days=1) + + +class TestWorkerRequests: + def test_avoid_worker_constraint(self): + Rota = generate_basic_rota() + + + Rota.add_shifts( + SingleShift( + sites=("group1",), name="a", length= 12.5, days=days, + workers_required=2, + force_as_block=False, + ), + ) + + Rota.add_worker_name_constraint_by_week(["worker1"], [1,2,3], ["a"]) + Rota.add_worker_name_constraint_by_week(["worker2"], [4,5,6], ["a"]) + Rota.add_worker_name_constraint_by_week(["worker3"], [7,8,9], ["a"]) + + Rota.build_and_solve(options={"ratio": 0.000}) + Rota.export_rota_to_html("avoid_grade_constraint") + + shift_summary = Rota.get_shift_summary_dict() + for worker_name in shift_summary: + worker_shifts = shift_summary[worker_name] + + worker = Rota.get_worker_by_name(worker_name) + + shift_string = Rota.get_worker_shift_list_string(worker) + + match worker_name: + case "worker1": + assert shift_string.startswith("-"*21) + case "worker2": + assert shift_string.startswith(21*"a"+"-"*21) + case "worker3": + assert shift_string.endswith("-"*21) + + + def test_avoid_worker_constraint_invalid_shift(self): + Rota = generate_basic_rota() + + Rota.add_shifts( + SingleShift( + sites=("group1",), name="a", length= 12.5, days=days, + workers_required=1, + force_as_block=False, + ), + ) + with pytest.raises(InvalidShift): + + Rota.add_worker_name_constraint_by_week(["worker1"], [1,2,3], ["b"]) + Rota.build_and_solve(options={"ratio": 0.000}) + + + Rota.add_shifts( + SingleShift( + sites=("group1",), name="b", length= 12.5, days=days, + workers_required=1, + force_as_block=False, + ), + ) + + Rota.add_worker_name_constraint_by_week([1], [1,2,3], ["b"]) + Rota.build_and_solve(options={"ratio": 0.000}) + assert Rota.results.solver.status == "ok" + + +# TODO: implement check for invalid workers +## def test_avoid_grade_constraint_invalid_worker(self): +## Rota = generate_basic_rota() +## +## Rota.add_shifts( +## SingleShift( +## sites=("group1",), name="a", length= 12.5, days=days, +## workers_required=2, +## force_as_block=False, +## ), +## ) +## with pytest.raises(ValueError): +## +## Rota.add_worker_name_constraint_by_week([4], [1,2,3], ["a"]) +## Rota.build_and_solve(options={"ratio": 0.000}) +## +## Rota.add_worker( +## Worker(name="worker4", site="group1", grade=4), +## ) +## +## +## Rota.add_worker_name_constraint_by_week([4], [1,2,3], ["a"]) +## Rota.build_and_solve(options={"ratio": 0.000}) +## assert Rota.results.solver.status == "ok" + + def test_avoid_worker_constraint_multiple_workers(self): + Rota = generate_basic_rota() + + Rota.add_shifts( + SingleShift( + sites=("group1",), name="a", length= 12.5, days=days, + workers_required=1, + force_as_block=False, + ), + ) + + Rota.add_worker_name_constraint_by_week(["worker1", "worker2"], [1,2,3], ["a"]) + Rota.build_and_solve(options={"ratio": 0.000}) + assert Rota.results.solver.status == "ok" + + Rota.add_worker_name_constraint_by_week(["worker3"], [1,2,3], ["a"]) + Rota.build_and_solve(options={"ratio": 0.000}) + assert Rota.results.solver.status == "warning" + + def test_avoid_worker_constraint_multiple_shifts(self): + Rota = generate_basic_rota() + + Rota.add_shifts( + SingleShift( + sites=("group1",), name="a", length= 12.5, days=days, + workers_required=1, + force_as_block=False, + ), + SingleShift( + sites=("group1",), name="b", length= 12.5, days=days, + workers_required=1, + force_as_block=False, + ), + ) + + Rota.add_worker_name_constraint_by_week(["worker1"], [1,2,3], ["a", "b"]) + Rota.add_worker_name_constraint_by_week(["worker2"], [4,5,6], ["a", "b"]) + Rota.add_worker_name_constraint_by_week(["worker3"], [7,8,9], ["a", "b"]) + + Rota.build_and_solve(options={"ratio": 0.000}) + Rota.export_rota_to_html("avoid_grade_constraint") + + shift_summary = Rota.get_shift_summary_dict() + for worker_name in shift_summary: + worker_shifts = shift_summary[worker_name] + + worker = Rota.get_worker_by_name(worker_name) + + shift_string = Rota.get_worker_shift_list_string(worker) + + match worker_name: + case "worker1": + assert shift_string.startswith("-"*21) + case "worker2": + assert "-"*21 in shift_string # hacky + case "worker3": + assert shift_string.endswith("-"*21) + + Rota.add_worker_name_constraint_by_week(["worker2"], [7,8,9], ["a"]) + Rota.build_and_solve(options={"ratio": 0.000}) + assert Rota.results.solver.status == "ok" + + Rota.add_worker_name_constraint_by_week(["worker1"], [7,8,9], ["a"]) + Rota.build_and_solve(options={"ratio": 0.000}) + assert Rota.results.solver.status == "warning" \ No newline at end of file