work towards paring of shifts

requirements.txt

@@ -14,3 +14,4 @@ requests
 loguru
 numpy
 highspy
+typer

rota/shifts.py

@@ -4,7 +4,7 @@ import itertools
 from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal
 import time
 
-from pydantic import BaseModel, Extra, constr
+from pydantic import BaseModel, ConfigDict, Extra, constr
 
 from pathlib import Path
 
@@ -81,6 +81,10 @@ class ShiftConstraint(BaseModel):
 class SingleShift(BaseModel):
     """Class to hold all details for a shift
 
+    start_date: The day that the shift starts on
+    end_date: The day that the shift ends on (the last day it should be rota'd for)
+
+    Due to week block assigments these (probably) need to be week ba
 
     Valid constraints
 
@@ -100,6 +104,9 @@ class SingleShift(BaseModel):
         options: (grade, min_number)
 
 
+
+
+
     """
 
     sites: List[str]
@@ -116,10 +123,13 @@ class SingleShift(BaseModel):
     hard_constrain_shift: bool = True
     bank_holidays_only: bool = False
     constraint: List[ShiftConstraint] = []
+    start_date: datetime.date | None = None
+    end_date: datetime.date | None = None
 
-    class Config:
+    model_config = ConfigDict(
-        extra = Extra.allow
+        extra = "allow",
-        orm_mode = True
+        #orm_mode = True
+    )
 
     def __init__(self, **data: Any):
         super().__init__(**data)
@@ -251,6 +261,8 @@ class RotaBuilder(object):
             "Worker/duplicate id",
             "Worker/duplicate name",
             "Worker/no valid shifts",
+            "Shift/invalid start date",
+            "Shift/invalid end date"
         ]
 
         self.results = None
@@ -266,6 +278,8 @@ class RotaBuilder(object):
 
         self.bank_holidays = bank_holidays
 
+        self.paired_shifts = []
+
     def set_rota_dates(self, start_date: datetime.date, weeks_to_rota: int):
         self.weeks_to_rota = weeks_to_rota
 
@@ -1100,6 +1114,16 @@ class RotaBuilder(object):
         for week in track(self.weeks, description="Generating week constraints..."):
             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):
                     self.model.constraints.add(
                         8
@@ -1215,7 +1239,7 @@ class RotaBuilder(object):
                     # Occurs if there are no shifts within a defined block
                     # TODO: test if this breaks (and we should check rathar than except)
                     self.add_warning(
-                        "max shifts per month constraint"
+                        "max shifts per month constraint",
                         f"Failed to constrain max_shifts_per_month for worker: {worker.name}"
                     )
                     pass
@@ -1702,6 +1726,7 @@ class RotaBuilder(object):
                 for constraint_shift in self.get_shifts_with_constraints(
                     "max_shifts_per_week"
                 ):
+                    try:
                         self.model.constraints.add(
                             constraint_shift.constraint_options["max_shifts_per_week"]
                             >= sum(
@@ -1714,6 +1739,9 @@ class RotaBuilder(object):
                                 if w == week
                             )
                         )
+                    except ValueError:
+                        # This happens if a shift if not assigned on the week (should we test for this instead?)
+                        pass
 
                 try:
                     self.model.constraints.add(
@@ -1902,7 +1930,7 @@ class RotaBuilder(object):
                 #        raise ValueError(f"Requires {shift.name} to have force_as_block")
 
                 # for shift in self.get_shifts_with_constraint("night"):
-                for shift in self.get_shifts():
+                for shift in self.get_shifts(week=week):
                     if shift.force_as_block:
                         # Force nights to be assigned in blocks
                         # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days)
@@ -2133,7 +2161,7 @@ class RotaBuilder(object):
                 # shift spans 7 days you may get >7 allocations in a row
                 # as it only checks for a different shift allocation
                 for constraint_shift in self.get_shifts_with_constraints(
-                    "pre",
+                    "pre", week=week
                 ):
                     for n in range(0, constraint_shift.constraint_options["pre"]):
                         if day in constraint_shift.days:
@@ -2181,7 +2209,7 @@ class RotaBuilder(object):
                 #        )
 
                 for constraint_shift in self.get_shifts_with_constraints(
-                    "post",
+                    "post", week=week
                 ):
                     for n in range(0, constraint_shift.constraint_options["post"]):
                         if day in constraint_shift.days:
@@ -2633,6 +2661,11 @@ class RotaBuilder(object):
         """
         return self.shifts_by_name[name].length
 
+    def pair_shifts(self, shift1: SingleShift, shift2: SingleShift):
+        # NOTE: currently designed to pair two shifts
+        # this could be extended...
+        self.pair_shifts.append(set(shift1, shift2))
+
     def build_shifts(self):
         """
         Process the added shifts
@@ -2655,6 +2688,24 @@ class RotaBuilder(object):
             if s.name in self.shift_names:
                 raise InvalidShift(f"Duplicate shift: {s.name}")
 
+            if s.start_date is not None:
+                if s.start_date > self.rota_end_date or s.start_date < self.start_date:
+                    self.add_warning(
+                        "Shift/invalid start date",
+                        f"Shift '{s.name}' has a start date outside of the rota limits ({s.start_date} vs [{self.start_date}---{self.rota_end_date}])",
+                    )
+            else:
+                s.start_date = self.start_date
+
+            if s.end_date is not None:
+                if s.end_date > self.rota_end_date or s.end_date < self.start_date:
+                    self.add_warning(
+                        "Shift/invalid end date",
+                        f"Shift '{s.name}' has an end date outside of the rota limits ({s.end_date} vs [{self.start_date}---{self.rota_end_date}])",
+                    )
+            else:
+                s.end_date = self.rota_end_date
+
             self.shifts_by_name[s.name] = s
             self.shift_names.append(s.name)
 
@@ -2667,10 +2718,15 @@ class RotaBuilder(object):
         self.shift_counts = defaultdict(int)
         for week, day in self.weeks_days_product:
             self.week_day_shifts_dict[(week, day)] = set()
+
+            current_date = self.week_day_date_map[(week, day)]
             for s in self.shifts:
+                # Check if the shift has started or ended
+                if current_date < s.start_date or current_date > s.end_date:
+                    continue
+
                 if s.bank_holidays_only:
-                    if self.week_day_date_map[(week, day)] in self.bank_holidays:
+                    if current_date in self.bank_holidays:
-                        print(s.name, self.week_day_date_map[(week, day)])
                         self.week_day_shift_product.append((week, day, s.name))
                         self.week_day_shiftclass_product.append((week, day, s))
                         self.week_day_shifts_dict[(week, day)].add(s.name)
@@ -2681,7 +2737,6 @@ class RotaBuilder(object):
                         self.week_day_shiftclass_product.append((week, day, s))
                         self.week_day_shifts_dict[(week, day)].add(s.name)
                         self.shift_counts[s.name] = self.shift_counts[s.name] + 1
-        # print(self.week_day_shift_product)
 
         self.max_pre = 1
         self.max_post = 1
@@ -2761,7 +2816,7 @@ class RotaBuilder(object):
         """
         return self.weeks_days_product
 
-    def get_week_day_combinations_for_shift(self, shift) -> list:
+    def get_week_day_combinations_for_shift(self, shift: SingleShift) -> list:
         return [
             (week, day)
             for week, day in self.get_week_day_combinations()
@@ -2779,12 +2834,24 @@ class RotaBuilder(object):
 
         return week_day_shifts
 
-    def get_shifts(self) -> List[SingleShift]:
+    def get_shifts(self, week: int | None=None) -> List[SingleShift]:
         """Returns a list of all the registered shifts
 
         Returns:
             List[SingleShift]: list of registered shifts (as SingleShift)
         """
+        if week is not None:
+            date = self.get_week_start_date(week)
+
+            shifts = []
+
+            for shift in self.shifts:
+                if shift.start_date < date <= shift.end_date:
+                    shifts.append(shift)
+
+            return shifts
+
+        else:
             return self.shifts
 
     def get_shifts_for_worker_site(self, worker_site):
@@ -2799,14 +2866,19 @@ class RotaBuilder(object):
     def get_shifts_with_constraint(self, constraint) -> List[SingleShift]:
         return [shift for shift in self.shifts if constraint in shift.constraints]
 
-    def get_shifts_with_constraints(self, *constraints) -> List[SingleShift]:
+    def get_shifts_with_constraints(self, *constraints, week : int | None =None) -> List[SingleShift]:
         shift_names = set()
 
         for constraint in constraints:
             shift_names.update(
-                [shift.name for shift in self.shifts if constraint in shift.constraints]
+                [
+                    shift.name
+                    for shift in self.get_shifts(week)
+                    if constraint in shift.constraints
+                ]
             )
 
+
         return [self.get_shift_by_name(s) for s in shift_names]
 
     def get_shift_names(self) -> List[ShiftName]:
@@ -2879,6 +2951,7 @@ class RotaBuilder(object):
             if (shift.force_as_block or shift.force_as_block_unless_nwd)
         ]
 
+    #def shifts_to_assign_or_force_as_blocks(self, week: int | None = None) -> List[str]:
     def shifts_to_assign_or_force_as_blocks(self) -> List[str]:
         s = self.shifts_to_assign_as_blocks()
         s.extend(self.shifts_to_force_as_blocks())
@@ -2923,7 +2996,6 @@ class RotaBuilder(object):
     def get_worker_details(self):
         w = defaultdict(list)
         for worker in self.workers:
-            # print(worker)
             w[worker.site].append(worker)
 
         l = []

rota/workers.py

@@ -1,7 +1,7 @@
 import datetime
 from collections import defaultdict
 from typing import Iterable, List, Literal, Optional
-from pydantic import BaseModel, Extra, validator
+from pydantic import BaseModel, ConfigDict, field_validator
 
 from rich.pretty import pprint
 from rota.console import console
@@ -37,7 +37,7 @@ class NonWorkingDays(BaseModel):
     start_date: datetime.date | None = None
     end_date: datetime.date | None = None
 
-    @validator("day")
+    @field_validator("day")
     def day_in(cls, day):
         for whole_day in whole_days:
             if day.lower() in whole_day:
@@ -83,8 +83,9 @@ class Worker(BaseModel):
     bank_holiday_extra: int = 0
     pair: int | str | None = None
 
-    class Config:
+    model_config = ConfigDict(
-        extra = Extra.allow
+        extra = "allow",
+    )
 
     #    def __init__(
     #        self,

test/test_shifts.py

@@ -1,13 +1,13 @@
 import datetime
 import pytest
-from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, WarningTermination, days
 
 from rota.workers import Worker
 
 import itertools
 
 
-def generate_basic_rota(weeks_to_rota=10):
+def generate_basic_rota(weeks_to_rota=10, workers=2):
     start_date = datetime.date(2022, 3, 7)
 
     Rota = RotaBuilder(
@@ -18,8 +18,8 @@ def generate_basic_rota(weeks_to_rota=10):
     # Add a few workers
     Rota.add_workers(
         [
-            Worker(name="worker1", site="group1", grade=1),
+            Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)
-            Worker(name="worker2", site="group1", grade=1),
+            #Worker(name="worker2", site="group1", grade=1),
         ]
     )
 
@@ -312,3 +312,130 @@ class TestShiftConstraints:
         Rota.export_rota_to_html("post")
 
         assert Rota.results.solver.status in ("warning", "error")
+
+
+class TestShiftDates:
+
+    def test_shift_start_date(self):
+        Rota = generate_basic_rota(workers=2)
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                start_date=(Rota.start_date + datetime.timedelta(days=7)),
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        #Rota.export_rota_to_html("test_shift_start_date")
+
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
+        assert Rota.get_workers_total_shifts()["worker1"] in (31, 32)
+        assert Rota.get_workers_total_shifts()["worker2"] in (31, 32)
+
+    def test_shift_start_date_end_date(self):
+        Rota = generate_basic_rota(workers=2)
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                start_date=(Rota.start_date + datetime.timedelta(days=7)),
+                end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        #Rota.export_rota_to_html("test_shift_start_date")
+
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
+        assert Rota.get_workers_total_shifts()["worker1"] in (17, 18)
+        assert Rota.get_workers_total_shifts()["worker2"] in (17, 18)
+
+        for worker in Rota.get_workers():
+            assert Rota.get_worker_shift_list_string(worker).startswith("-"*7)
+            assert Rota.get_worker_shift_list_string(worker).endswith("-"*28)
+
+    def test_shift_invalid_start_date(self):
+        Rota = generate_basic_rota(workers=2)
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                start_date=(Rota.start_date - datetime.timedelta(days=6)),
+            ),
+        )
+
+        with pytest.raises(WarningTermination):
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+    def test_shift_invalid_start_date2(self):
+        Rota = generate_basic_rota(workers=2)
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                start_date=(Rota.rota_end_date + datetime.timedelta(days=1)),
+            ),
+        )
+
+        with pytest.raises(WarningTermination):
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+    def test_shift_invalid_end_date(self):
+        Rota = generate_basic_rota(workers=2)
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                end_date=(Rota.rota_end_date + datetime.timedelta(days=1)),
+            ),
+        )
+
+        with pytest.raises(WarningTermination):
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+    def test_shift_invalid_end_date2(self):
+        Rota = generate_basic_rota(workers=2)
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                end_date=(Rota.start_date - datetime.timedelta(days=1)),
+            ),
+        )
+
+        with pytest.raises(WarningTermination):
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+
+    def test_shift_start_date_end_date_block(self):
+        Rota = generate_basic_rota(workers=2)
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                assign_as_block=True,
+                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                start_date=(Rota.start_date + datetime.timedelta(days=7)),
+                end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        #Rota.export_rota_to_html("test_shift_start_date")
+
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
+        assert Rota.get_workers_total_shifts()["worker1"] in (17, 18)
+        assert Rota.get_workers_total_shifts()["worker2"] in (17, 18)
+
+        for worker in Rota.get_workers():
+            assert Rota.get_worker_shift_list_string(worker).startswith("-"*7)
+            assert Rota.get_worker_shift_list_string(worker).endswith("-"*28)