fix bug with pre / post shift constraints

rota/shifts.py

@@ -2626,10 +2626,9 @@ class RotaBuilder(object):
                 ):
                     for n in range(0, constraint_shift.constraint_options["pre"]):
                         if day in constraint_shift.days:
-                            # week_start_date = self.get_week_start_date(week)
+                            # Only apply if worker can work this shift
-                            # try:
+                            if worker.site not in constraint_shift.sites:
-                            # if constraint_shift.start_date <= week_start_date < constraint_shift.end_date:
+                                continue
-                            #    works = 1
                             try:
                                 works = self.model.works[
                                     worker.id, week, day, constraint_shift.name
@@ -2649,50 +2648,23 @@ class RotaBuilder(object):
                                     )
                                     if shiftname != constraint_shift.name
                                     for w in workers
+                                    if w.site in constraint_shift.sites  # Only workers who can work this shift
                                 )
                             )
-                # for constraint_shift in self.get_shifts_with_constraint( "pre"):
-                #    to_sum = []
-                #    for n in range(
-                #        0, constraint_shift.constraint_options["pre"]
-                #    ):
-                #        for w in workers:
-                #            for wk, dy, shiftname in self.get_week_day_shift_names_by_week_day(
-                #                pre_map[n][1], pre_map[n][2]
-                #            ):
-                #                if shiftname != constraint_shift.name:
-                #                    to_sum.append(pre_map[n][0]
-                #                    * self.model.works[
-                #                        w.id, wk, dy, shiftname
-                #                    ])
-
-                #    if day in constraint_shift.days:
-                #        self.model.constraints.add(
-                #            1
-                #            >= self.model.works[
-                #                worker.id, week, day, constraint_shift.name
-                #            ]
-                #            + sum(
-                #                to_sum
-                #            )
-                #        )
 
                 for constraint_shift in self.get_shifts_with_constraints(
                     "post"  # , week=week
                 ):
                     for n in range(0, constraint_shift.constraint_options["post"]):
                         if day in constraint_shift.days:
-                            # week_start_date = self.get_week_start_date(week)
+                            # Only apply if worker can work this shift
-                            # if constraint_shift.start_date <= week_start_date < constraint_shift.end_date:
+                            if worker.site not in constraint_shift.sites:
-                            # works = 1
+                                continue
-                            # else:
                             try:
                                 works = self.model.works[
                                     worker.id, week, day, constraint_shift.name
                                 ]
                             except KeyError:
-                                # If we can't find the key the shift is not assigned
-                                # so we don't need to worry about it
                                 continue
                             self.model.constraints.add(
                                 1
@@ -2707,10 +2679,9 @@ class RotaBuilder(object):
                                     )
                                     if shiftname != constraint_shift.name
                                     for w in workers
+                                    if w.site in constraint_shift.sites  # Only workers who can work this shift
                                 )
                             )
-                            # except KeyError:
-                            #    pass
 
                 # Night constraint means we won't assign a shift the day before
                 # an unavailability

test/test_rota_clear.py

@@ -3,6 +3,27 @@ from rota.shifts import RotaBuilder, SingleShift, days
 from rota.workers import Worker
 import datetime
 
+
+def generate_basic_rota(
+    weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7)
+):
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+
+    # Add a few workers
+    if workers:
+        Rota.add_workers(
+            [
+                Worker(name=f"worker{i:02}", site="group1", grade=1)
+                for i in range(1, workers + 1)
+            ]
+        )
+
+    return Rota
+
+
 @pytest.fixture
 def demo_rota_clear():
     weeks_to_rota = 10
@@ -28,6 +49,7 @@ def demo_rota_clear():
     Rota.add_workers(workers)
     return Rota
 
+
 def test_pre(demo_rota_clear):
     Rota = demo_rota_clear
     Rota.add_shifts(
@@ -83,3 +105,350 @@ def test_pre(demo_rota_clear):
         # assert "c-a" not in shifts_string
         # assert "b-c" not in shifts_string
         # assert "c-b" not in shifts_string
+
+
+def test_pre_post_clear(demo_rota_clear):
+    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[:5],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 1,
+                },
+                {"name": "post", "options": 1},
+            ],
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[5:],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 1,
+                },
+                {"name": "post", "options": 1},
+            ],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_pre_post_clear", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+    for worker in Rota.workers:
+        shifts = Rota.get_worker_shift_list(worker)
+        shifts_string = "".join([i if i != "" else "-" for i in shifts])
+
+        # check for patterns that should not occur
+        assert "ab" not in shifts_string
+        assert "ba" not in shifts_string
+        # assert "a-b" not in shifts_string
+        # assert "b-a" not in shifts_string
+
+
+def test_pre_post_clear2(demo_rota_clear):
+    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[:5],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 2,
+                },
+                {"name": "post", "options": 2},
+            ],
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[5:],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 2,
+                },
+                {"name": "post", "options": 2},
+            ],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_pre_post_clear", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+    for worker in Rota.workers:
+        shifts = Rota.get_worker_shift_list(worker)
+        shifts_string = "".join([i if i != "" else "-" for i in shifts])
+
+        # check for patterns that should not occur
+        assert "ab" not in shifts_string
+        assert "ba" not in shifts_string
+        assert "a-b" not in shifts_string
+        assert "b-a" not in shifts_string
+
+
+def test_pre_post_clear_multi_group(demo_rota_clear):
+    Rota = generate_basic_rota(workers=2, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    Rota.add_workers(
+        (
+            Worker(name="worker3", site="group2", grade=1),
+            Worker(name="worker4", site="group2", grade=1),
+        )
+    )
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 2,
+                },
+                {"name": "post", "options": 2},
+            ],
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[5:],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 2,
+                },
+                {"name": "post", "options": 2},
+            ],
+        ),
+        SingleShift(
+            sites=("group2",),
+            name="c",
+            length=8,
+            days=days[5:],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 2,
+                },
+                {"name": "post", "options": 2},
+            ],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_pre_post_clear", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+    for worker in Rota.workers:
+        shifts = Rota.get_worker_shift_list(worker)
+        shifts_string = "".join([i if i != "" else "-" for i in shifts])
+
+        # check for patterns that should not occur
+        assert "ab" not in shifts_string
+        assert "ba" not in shifts_string
+        assert "a-b" not in shifts_string
+        assert "b-a" not in shifts_string
+
+        assert "ac" not in shifts_string
+        assert "ca" not in shifts_string
+        assert "a-c" not in shifts_string
+        assert "c-a" not in shifts_string
+
+def test_pre_post_clear_force_assign(demo_rota_clear):
+    Rota = generate_basic_rota(workers=6, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_force_assign_with("a", "o")
+
+    Rota.add_workers(
+        (
+            Worker(name="workera", site="group2", grade=1),
+            Worker(name="workerb", site="group2", grade=1),
+        )
+    )
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 1,
+                },
+                {"name": "post", "options": 1},
+            ],
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="o",
+            length=8,
+            days=days[:7],
+            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
+            #constraint=[
+            #    {
+            #        "name": "pre",
+            #        "options": 2,
+            #    },
+            #    {"name": "post", "options": 2},
+            #],
+        ),
+        SingleShift(
+            sites=("group2",),
+            name="c",
+            length=8,
+            days=days[5:],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+            #constraint=[
+            #    {
+            #        "name": "pre",
+            #        "options": 1,
+            #    },
+            #    {"name": "post", "options": 1},
+            #],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_pre_post_clear", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+
+def test_pre_post_clear_force_assign2(demo_rota_clear):
+    Rota = generate_basic_rota(workers=6, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_force_assign_with("a", "o")
+        worker.add_force_assign_with("b", "o")
+
+    Rota.add_workers(
+        (
+            Worker(name="workera", site="group2", grade=1),
+            Worker(name="workerb", site="group2", grade=1),
+        )
+    )
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                #{
+                #    "name": "pre",
+                #    "options": 1,
+                #},
+                #{"name": "post", "options": 1},
+            ],
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="o",
+            length=8,
+            days=days[:7],
+            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
+            constraint=[
+                #{
+                #    "name": "pre",
+                #    "options": 2,
+                #},
+                #{"name": "post", "options": 2},
+            ],
+        ),
+        SingleShift(
+            sites=("group2",),
+            name="c",
+            length=8,
+            days=days[5:],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+            constraint=[
+                {
+                    "name": "pre",
+                    "options": 1,
+                },
+                {"name": "post", "options": 1},
+            ],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_pre_post_clear", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+    for worker in Rota.workers:
+        shifts = Rota.get_worker_shift_list(worker)
+        shifts_string = "".join([i if i != "" else "-" for i in shifts])
+
+        # check for patterns that should not occur
+        assert "ac" not in shifts_string
+        assert "ca" not in shifts_string

test/test_workers.py

@@ -964,10 +964,11 @@ def test_worker_assign_as_block_preference_multiple_shifts():
         ),
     )
 
-    shift_list_1 = Rota.get_worker_shift_list_string(worker1)
-    assert shift_list_1.count("aaaa") == 3, "Worker 1 should have all 'a' shifts grouped together (block)"
 
     Rota.build_and_solve(options={"ratio": 0.0})
+
+    shift_list_1 = Rota.get_worker_shift_list_string(worker1)
+    assert shift_list_1.count("aaaa") == 3, "Worker 1 should have all 'a' shifts grouped together (block)"
     Rota.export_rota_to_html("test_worker_assign_as_block_preference", folder="tests")
 
 def test_worker_assign_as_block_preference_multiple_sites():