add worker requirements

gen.py

@@ -2,7 +2,7 @@ import datetime
 import os
 import sys
 import time
-from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
+from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WorkerRequirement, days
 import typer
 import subprocess
 
@@ -30,6 +30,7 @@ from rota.workers import (
     OutOfProgramme,
 )
 
+
 @app.command()
 def main(
     suspend_on_finish: bool = False,
@@ -40,11 +41,13 @@ def main(
     weeks: int = 26,
     bom: int = 1,
 ):
-
     rota_start_date = start_date.date()
 
     Rota = RotaBuilder(
-        rota_start_date, weeks_to_rota=weeks, balance_offset_modifier=bom, name="proc_rota"
+        rota_start_date,
+        weeks_to_rota=weeks,
+        balance_offset_modifier=bom,
+        name="proc_rota",
     )
     # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
     Rota.constraint_options["balance_weekends"] = True
@@ -54,6 +57,17 @@ def main(
     Rota.constraint_options["hard_constrain_pair_separation"] = True
     # Rota.constraint_options["avoid_st2_first_month"] = True
 
+    wr = [
+        WorkerRequirement(
+            end_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
+            number=3,
+        ),
+        WorkerRequirement(
+            start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
+            number=4,
+        ),
+    ]
+
     Rota.add_shifts(
         SingleShift(
             sites=(
@@ -137,8 +151,8 @@ def main(
             length=12.5,
             days=days[4:],
             balance_offset=3,
-            #rota_on_nwds=True,
+            # rota_on_nwds=True,
-            #force_as_block=True,
+            # force_as_block=True,
             assign_as_block=True,
             constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
             # force_as_block_unless_nwd=True
@@ -209,7 +223,7 @@ def main(
             balance_offset=3.9,
             balance_weighting=1,
             # hard_constrain_shift=False,
-            workers_required=3,
+            workers_required=wr,
             force_as_block=True,
             rota_on_nwds=True,
             constraint=[
@@ -223,7 +237,7 @@ def main(
                 {"name": "limit_grade_number", "options": {2: 1}},
                 {"name": "minimum_grade_number", "options": (4, 1)},
             ],
-            end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
+            #end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
         ),
         SingleShift(
             sites=[
@@ -239,7 +253,7 @@ def main(
             balance_offset=2.9,
             balance_weighting=1,
             # hard_constrain_shift=False,
-            workers_required=3,
+            workers_required=wr,
             force_as_block=True,
             rota_on_nwds=True,
             constraint=[
@@ -253,62 +267,62 @@ def main(
                 {"name": "limit_grade_number", "options": {2: 1}},
                 {"name": "minimum_grade_number", "options": (4, 1)},
             ],
-            end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
+            #end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
-        ),
-        SingleShift(
-            sites=sites,
-            name="night_weekday4",
-            length=12.25,
-            days=days[:4],
-            balance_offset=3.9,
-            balance_weighting=1,
-            # hard_constrain_shift=False,
-            workers_required=NIGHT_REGISTRAR_NUMBER,
-            force_as_block=True,
-            rota_on_nwds=True,
-            constraint=[
-                {"name": "night"},
-                {"name": "pre", "options": 2},
-                {"name": "post", "options": 2},
-                {
-                    "name": "require_remote_site_presence_week",
-                    "options": ("plymouth", 1),
-                },
-                {"name": "limit_grade_number", "options": {2: 1}},
-                {"name": "minimum_grade_number", "options": (4, 1)},
-            ],
-            start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
-        ),
-        SingleShift(
-            sites=[
-                *sites,
-                "plymouth twilights, weekends and weekend nights",
-                "truro twilights and weekend nights",
-                "weekend nights",
-                "truro twilights, weekends and weekend nights",
-            ],
-            name="night_weekend4",
-            length=12.25,
-            days=days[4:],
-            balance_offset=2.9,
-            balance_weighting=1,
-            # hard_constrain_shift=False,
-            workers_required=NIGHT_REGISTRAR_NUMBER,
-            force_as_block=True,
-            rota_on_nwds=True,
-            constraint=[
-                {"name": "night"},
-                {"name": "pre", "options": 2},
-                {"name": "post", "options": 3},
-                {
-                    "name": "require_remote_site_presence_week",
-                    "options": ("plymouth", 1),
-                },
-                {"name": "limit_grade_number", "options": {2: 1}},
-                {"name": "minimum_grade_number", "options": (4, 1)},
-            ],
-            start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
         ),
+        # SingleShift(
+        #    sites=sites,
+        #    name="night_weekday4",
+        #    length=12.25,
+        #    days=days[:4],
+        #    balance_offset=3.9,
+        #    balance_weighting=1,
+        #    # hard_constrain_shift=False,
+        #    workers_required=NIGHT_REGISTRAR_NUMBER,
+        #    force_as_block=True,
+        #    rota_on_nwds=True,
+        #    constraint=[
+        #        {"name": "night"},
+        #        {"name": "pre", "options": 2},
+        #        {"name": "post", "options": 2},
+        #        {
+        #            "name": "require_remote_site_presence_week",
+        #            "options": ("plymouth", 1),
+        #        },
+        #        {"name": "limit_grade_number", "options": {2: 1}},
+        #        {"name": "minimum_grade_number", "options": (4, 1)},
+        #    ],
+        #    start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
+        # ),
+        # SingleShift(
+        #    sites=[
+        #        *sites,
+        #        "plymouth twilights, weekends and weekend nights",
+        #        "truro twilights and weekend nights",
+        #        "weekend nights",
+        #        "truro twilights, weekends and weekend nights",
+        #    ],
+        #    name="night_weekend4",
+        #    length=12.25,
+        #    days=days[4:],
+        #    balance_offset=2.9,
+        #    balance_weighting=1,
+        #    # hard_constrain_shift=False,
+        #    workers_required=NIGHT_REGISTRAR_NUMBER,
+        #    force_as_block=True,
+        #    rota_on_nwds=True,
+        #    constraint=[
+        #        {"name": "night"},
+        #        {"name": "pre", "options": 2},
+        #        {"name": "post", "options": 3},
+        #        {
+        #            "name": "require_remote_site_presence_week",
+        #            "options": ("plymouth", 1),
+        #        },
+        #        {"name": "limit_grade_number", "options": {2: 1}},
+        #        {"name": "minimum_grade_number", "options": (4, 1)},
+        #    ],
+        #    start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
+        # ),
     )
 
     # Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"])
@@ -381,7 +395,7 @@ def main(
                     elif "to" in dates:
                         s, e = dates.split(" to ")
                     else:
-                        raise ValueError(f"Cannot parse OOP dates: {dates}")
+                        raise ValueError(f"Cannot parse OOP dates: '{oop}' for {worker}")
                     try:
                         formatted_oops.append(
                             {
@@ -474,13 +488,13 @@ def main(
     # Rota.build_model()
 
     solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
-    #solver_options = {"seconds": time_to_run, "threads": 10}
+    # solver_options = {"seconds": time_to_run, "threads": 10}
 
     # start_time = time.time()
     Rota.build_and_solve(solver_options, export=True, solve=solve, solver="appsi_highs")
 
-    #Rota.solve_shifts_by_block(solver_options, block_length=13)
+    # Rota.solve_shifts_by_block(solver_options, block_length=13)
-    #Rota.solve_shifts_individually(solver_options)
+    # Rota.solve_shifts_individually(solver_options)
     # Rota.solve_model(options=solver_options)
     # end_time = time.time()
 
@@ -503,7 +517,9 @@ def main(
 
     # Rota.export_rota_to_html("proc_rota")
     # Rota.export_rota_to_csv("rota")
-    subprocess.run(["scp", Rota.exported_rota_file, "ross@46.101.13.46:proc/proc-rota/output/"])
+    subprocess.run(
+        ["scp", Rota.exported_rota_file, "ross@46.101.13.46:proc/proc-rota/output/"]
+    )
 
     if suspend_on_finish:
         os.system("systemctl suspend")

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, ConfigDict, Extra, constr
+from pydantic import BaseModel, ConfigDict, Extra, constr, model_validator, StrictInt
 
 from pathlib import Path
 
@@ -78,6 +78,23 @@ class ShiftConstraint(BaseModel):
     options: bool | int | Dict | tuple = False
 
 
+class WorkerRequirement(BaseModel):
+    """
+    start date is inclusive, end date is not
+    """
+    # NOTE: if number is changed mid week block allocations may not work?
+    number: int = 1
+    start_date: datetime.date | None = None
+    end_date: datetime.date | None = None
+
+    @model_validator(mode="after")
+    def check_dates(self):
+        if self.start_date is not None and self.end_date is not None:
+            if self.start_date > self.end_date:
+                raise ValueError("Start date must be before end date")
+        return self
+
+
 class SingleShift(BaseModel):
     """Class to hold all details for a shift
 
@@ -112,24 +129,24 @@ class SingleShift(BaseModel):
     sites: List[str]
     name: str
     length: float
-    days: str | List[str]
+    days: str | list[str]
     balance_offset: float | None = None
     balance_weighting: float = 1
-    workers_required: float = 1
+    workers_required: StrictInt | list[WorkerRequirement] = 1
     rota_on_nwds: bool = False
     assign_as_block: bool = False
     force_as_block: bool = False
     force_as_block_unless_nwd: bool = False
     hard_constrain_shift: bool = True
     bank_holidays_only: bool = False
-    constraint: List[ShiftConstraint] = []
+    constraint: list[ShiftConstraint] = []
     start_date: datetime.date | None = None
     end_date: datetime.date | None = None
     pair_proxy: bool = False
 
     model_config = ConfigDict(
-        extra = "allow",
+        extra="allow",
-        #orm_mode = True
+        # orm_mode = True
     )
 
     def __init__(self, **data: Any):
@@ -182,6 +199,21 @@ class SingleShift(BaseModel):
     def get_shift_number(self):
         return len(self.days)
 
+    def get_worker_requirement_by_date(self, date: datetime.date):
+        """Must be called after build_shifts
+        
+        Returns the number of workers required for a given date
+        
+        If none are required 0 is returned"""
+        if isinstance(self.workers_required, int):
+            return self.workers_required
+
+        for wr in self.workers_required:
+            if wr.start_date <= date < wr.end_date:
+                return wr.number
+
+        return 0
+
 
 class RotaBuilder(object):
     """Class to hold and manipulate shifts"""
@@ -263,7 +295,7 @@ class RotaBuilder(object):
             "Worker/duplicate name",
             "Worker/no valid shifts",
             "Shift/invalid start date",
-            "Shift/invalid end date"
+            "Shift/invalid end date",
         ]
 
         self.results = None
@@ -370,8 +402,7 @@ class RotaBuilder(object):
 
         if (
             not results.solver.termination_condition == TerminationCondition.infeasible
-            
+            # and not results.solver.status == SolverStatus.aborted
-            #and not results.solver.status == SolverStatus.aborted
         ):
             self.model.solutions.load_from(results)
 
@@ -1119,12 +1150,17 @@ class RotaBuilder(object):
                 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)
+                print(
+                    week,
+                    self.get_week_start_date(week),
+                    shift.start_date,
+                    shift.end_date,
+                )
 
-                #if self.get_week_start_date(week) <= shift.start_date:
+                # if self.get_week_start_date(week) <= shift.start_date:
                 #    continue
 
-                #if self.get_week_start_date(week) > shift.end_date:
+                # if self.get_week_start_date(week) > shift.end_date:
                 #    continue
 
                 for worker in self.get_workers_for_shift(shift):
@@ -1151,6 +1187,10 @@ class RotaBuilder(object):
                     )
                 )
 
+                workers_required = shift.get_worker_requirement_by_date(
+                    self.get_week_start_date(week)
+                )
+
                 if shift.force_as_block_unless_nwd:
                     workers = self.get_workers_for_shift(shift)
                     # Get workers who have a nwd on the shift
@@ -1188,7 +1228,7 @@ class RotaBuilder(object):
                                 ]
                                 for worker in self.workers
                             )
-                            <= shift.workers_required
+                            <= workers_required
                         )
 
                     else:
@@ -1199,19 +1239,20 @@ class RotaBuilder(object):
                                 ]
                                 for worker in full_workers
                             )
-                            <= shift.workers_required + 1
+                            <= workers_required + 1
                         )
 
                 elif shift.force_as_block:
-                    self.model.constraints.add(
+                    if workers_required:
-                        sum(
+                        self.model.constraints.add(
-                            self.model.blocks_worker_shift_assigned[
+                            sum(
-                                worker.id, week, shift.name
+                                self.model.blocks_worker_shift_assigned[
-                            ]
+                                    worker.id, week, shift.name
-                            for worker in self.workers
+                                ]
-                        ) 
+                                for worker in self.workers
-                        <= shift.workers_required
+                            )
-                    )
+                            <= workers_required
+                        )
 
         # Most of our constraints apply per worker
         worker: Worker
@@ -1246,7 +1287,7 @@ class RotaBuilder(object):
                     # TODO: test if this breaks (and we should check rathar than except)
                     self.add_warning(
                         "max shifts per month constraint",
-                        f"Failed to constrain max_shifts_per_month for worker: {worker.name}"
+                        f"Failed to constrain max_shifts_per_month for worker: {worker.name}",
                     )
                     pass
 
@@ -1326,8 +1367,12 @@ class RotaBuilder(object):
                 ):  # Each site specfies which sites self.workers can fullfill it
                     # calaculate the total number of shifts that need assigning over the rota
                     # total_shifts = (len(self.weeks) * len(shift.days) *
+                    # TODO: check if shift_counts is still required?
+                    #total_shifts = (
+                    #    self.shift_counts[shift.name] * shift.workers_required
+                    #)
                     total_shifts = (
-                        self.shift_counts[shift.name] * shift.workers_required
+                        self.shift_worker_counts[shift.name]
                     )
 
                     full_time_equivalent_joined = sum(
@@ -1736,7 +1781,9 @@ class RotaBuilder(object):
                         self.model.constraints.add(
                             constraint_shift.constraint_options["max_shifts_per_week"]
                             >= sum(
-                                self.model.works[worker.id, w, day, constraint_shift.name]
+                                self.model.works[
+                                    worker.id, w, day, constraint_shift.name
+                                ]
                                 # for shiftname in self.get_shift_names_by_week_day(week, day)
                                 # for day in self.days
                                 for w, day in self.get_week_day_combinations_for_shift(
@@ -1938,21 +1985,24 @@ class RotaBuilder(object):
                 # for shift in self.get_shifts_with_constraint("night"):
                 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)
-                        #     )
 
-                        # Force night shifts to be assigned in blocks
+                        if shift.get_worker_requirement_by_date(
-                        self.model.constraints.add(
+                            self.get_week_start_date(week)
-                            self.model.shift_week_worker_assigned[
+                        ):
-                                shift.name, week, worker.id
+                            # Force shifts to be assigned in blocks
-                            ]
+                            try:
-                            * len(shift.days)
+                                self.model.constraints.add(
-                            == sum(
+                                    self.model.shift_week_worker_assigned[
-                                self.model.works[worker.id, week, day, shift.name]
+                                        shift.name, week, worker.id
-                                for day in shift.days
+                                    ]
-                            )
+                                    * len(shift.days)
-                        )
+                                    == sum(
+                                        self.model.works[worker.id, week, day, shift.name]
+                                        for day in shift.days
+                                    )
+                                )
+                            except KeyError:
+                                pass # blocks with workerd requirements
 
                 #
 
@@ -2167,13 +2217,13 @@ 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"#, week=week
+                    "pre"  # , week=week
                 ):
                     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)
+                            # week_start_date = self.get_week_start_date(week)
-                            #try:
+                            # try:
-                            #if constraint_shift.start_date <= week_start_date < constraint_shift.end_date:
+                            # if constraint_shift.start_date <= week_start_date < constraint_shift.end_date:
                             #    works = 1
                             try:
                                 works = self.model.works[
@@ -2223,14 +2273,14 @@ class RotaBuilder(object):
                 #        )
 
                 for constraint_shift in self.get_shifts_with_constraints(
-                    "post"#, week=week
+                    "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)
+                            # week_start_date = self.get_week_start_date(week)
-                            #if constraint_shift.start_date <= week_start_date < constraint_shift.end_date:
+                            # if constraint_shift.start_date <= week_start_date < constraint_shift.end_date:
-                                #works = 1
+                            # works = 1
-                            #else:
+                            # else:
                             try:
                                 works = self.model.works[
                                     worker.id, week, day, constraint_shift.name
@@ -2254,7 +2304,7 @@ class RotaBuilder(object):
                                     for w in workers
                                 )
                             )
-                            #except KeyError:
+                            # except KeyError:
                             #    pass
 
                 #                for constraint_shift in self.get_shifts_with_constraints(
@@ -2688,10 +2738,15 @@ class RotaBuilder(object):
         """
         return self.shifts_by_name[name].length
 
-    def pair_shifts(self, shift1: SingleShift, shift2: SingleShift):
+    def get_date_by_week_day(self, week: WeekInt, day: DayStr) -> datetime.date:
+        return self.week_day_date_map[(week, day)]
+
+    def pair_shifts(self, *shifts: SingleShift):
         # NOTE: currently designed to pair two shifts
         # this could be extended...
-        self.paired_shifts.append([shift1, shift2])
+        if len(shifts) < 2:
+            raise ValueError(f"Must pair at least two shifts: {shifts}")
+        self.paired_shifts.append([*shifts])
 
     def build_shifts(self):
         """
@@ -2711,11 +2766,6 @@ class RotaBuilder(object):
         for s in self.shifts:
             pass
 
-        ## For paired shifts we create proxy shifts
-        #self.paired_shift_map = {}
-        #for shifts in self.paired_shifts:
-
-
         for s in self.shifts:
             if s.name in self.shift_names:
                 raise InvalidShift(f"Duplicate shift: {s.name}")
@@ -2738,6 +2788,33 @@ class RotaBuilder(object):
             else:
                 s.end_date = self.rota_end_date
 
+            # Check worker requirements
+            if not isinstance(s.workers_required, int):
+                # Validate the worker requirements
+                print(s.workers_required)
+                for worker_requirement in s.workers_required:
+                    if not isinstance(worker_requirement, WorkerRequirement):
+                        raise InvalidShift(
+                            f"Invalid worker requirement: {worker_requirement}"
+                        )
+                    
+                    if worker_requirement.start_date is None:
+                        worker_requirement.start_date = s.start_date
+                    elif worker_requirement.start_date > s.end_date:
+                        self.add_warning(
+                            "Shift/worker requirement start date",
+                            f"Worker requirement start date is after shift end date: {worker_requirement.start_date} > {s.end_date}"
+                        )
+                    
+                    if worker_requirement.end_date is None:
+                        worker_requirement.end_date = s.end_date
+                    elif worker_requirement.end_date < s.start_date:
+                        self.add_warning(
+                            "Shift/worker requirement end date",
+                            f"Worker requirement end date is before shift start date: {worker_requirement.end_date} < {s.start_date}"
+                        )
+
+
             self.shifts_by_name[s.name] = s
             self.shift_names.append(s.name)
 
@@ -2748,27 +2825,39 @@ class RotaBuilder(object):
                 self.sites.add(site)
 
         self.shift_counts = defaultdict(int)
+        self.shift_worker_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)]
+            current_date = self.get_date_by_week_day(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
 
+                # Shift not required if there are no worker requirements
+                workers_required = s.get_worker_requirement_by_date(current_date)
+                if not workers_required:
+                    self.add_warning(
+                        "Shift/worker requirement is zero",
+                        f"Shift '{s.name}' has no worker requirements on {current_date}",
+                    )
+                    continue
+
                 if s.bank_holidays_only:
                     if current_date in self.bank_holidays:
                         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)
                         self.shift_counts[s.name] = self.shift_counts[s.name] + 1
+                        self.shift_worker_counts[s.name] = self.shift_worker_counts[s.name] + workers_required
                 else:
                     if day in s.days:
                         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)
                         self.shift_counts[s.name] = self.shift_counts[s.name] + 1
+                        self.shift_worker_counts[s.name] = self.shift_worker_counts[s.name] + workers_required
 
         self.max_pre = 1
         self.max_post = 1
@@ -2778,6 +2867,19 @@ class RotaBuilder(object):
         for shift in self.get_shifts_with_constraint("post"):
             self.max_post = max(shift.constraint_options["post"], self.max_post)
 
+        ## For paired shifts we create proxy shifts
+        # self.paired_shift_map = {}
+        # for shifts in self.paired_shifts:
+        #    sites = shifts[0].sites
+        #    for shift in shifts:
+        #        self.paired_shift_map[shift] = shifts
+        #
+        #        if shift.sites != sites:
+        #            # This will lead to odd/broken behaviour
+        #            self.add_warning(
+        #                "Paired shifts/site mismatch",
+        #                f"Shift '{shift.name}' has different sites to paired shifts",
+        #            )
 
     #        # Todo replace with week_day.....
     #        self.day_shift_product = []
@@ -2867,7 +2969,7 @@ class RotaBuilder(object):
 
         return week_day_shifts
 
-    def get_shifts(self, week: int | None=None) -> List[SingleShift]:
+    def get_shifts(self, week: int | None = None) -> List[SingleShift]:
         """Returns a list of all the registered shifts
 
         Returns:
@@ -2880,7 +2982,8 @@ class RotaBuilder(object):
 
             for shift in self.shifts:
                 if shift.start_date < date <= shift.end_date:
-                    shifts.append(shift)
+                    if shift.get_worker_requirement_by_date(date):
+                        shifts.append(shift)
 
             return shifts
 
@@ -2899,7 +3002,9 @@ 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, week : int | None =None) -> List[SingleShift]:
+    def get_shifts_with_constraints(
+        self, *constraints, week: int | None = None
+    ) -> List[SingleShift]:
         shift_names = set()
 
         for constraint in constraints:
@@ -2911,7 +3016,6 @@ class RotaBuilder(object):
                 ]
             )
 
-
         return [self.get_shift_by_name(s) for s in shift_names]
 
     def get_shift_names(self) -> List[ShiftName]:
@@ -2952,8 +3056,11 @@ class RotaBuilder(object):
         """
         l = []
         for week, day, shift in self.week_day_shiftclass_product:
+            workers_required = shift.get_worker_requirement_by_date(
+                self.get_date_by_week_day(week, day)
+            )
             # if day in shift.days:
-            l.append((week, day, shift.name, shift.workers_required, shift.sites))
+            l.append((week, day, shift.name, workers_required, shift.sites))
 
         return l
 
@@ -2984,7 +3091,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, 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())

rota/workers.py

@@ -32,6 +32,27 @@ class NotAvailableToWork(BaseModel):
     reason: str = "unknown"
 
 
+def generate_not_available_to_works(
+    start_date: datetime.date,
+    end_date: datetime.date | None = None,
+    days: int | None = None,
+    reason: str = "unknown",
+) -> List[NotAvailableToWork]:
+    if end_date is None and days is None:
+        raise ValueError("Either end_date or days must be provided")
+    if end_date is not None and days is not None:
+        raise ValueError("Only one of end_date or days must be provided")
+
+    if end_date is not None:
+        days = (end_date - start_date).days
+    
+    not_available_to_works = []
+    for day in range(days):
+        not_available_to_works.append(NotAvailableToWork(date=start_date + datetime.timedelta(days=day), reason=reason))
+
+    return not_available_to_works
+
+
 class NonWorkingDays(BaseModel):
     day: days
     start_date: datetime.date | None = None
@@ -68,7 +89,7 @@ class Worker(BaseModel):
     # rules is easier.
     grade: int
     # We can either have a user generated ID
-    id: Optional[int| uuid.UUID| str] = None
+    id: Optional[int | uuid.UUID | str] = None
     fte: int = 100
     nwds: list[NonWorkingDays] = []
     start_date: datetime.date | None = None
@@ -84,7 +105,7 @@ class Worker(BaseModel):
     pair: int | str | None = None
 
     model_config = ConfigDict(
-        extra = "allow",
+        extra="allow",
     )
 
     #    def __init__(
@@ -111,14 +132,17 @@ class Worker(BaseModel):
 
     def load_rota(self, Rota: "RotaBuilder"):
         self.proportion_rota_to_work: float = 1
-        self.shift_target_number: dict[ str, int ] = defaultdict(int)
+        self.shift_target_number: dict[str, int] = defaultdict(int)
 
         if self.id is None:
             self.id = uuid.uuid4()
 
         if self.start_date is not None and self.end_date is not None:
             if self.start_date >= self.end_date:
-                Rota.add_warning("Worker/Early end date", f"{self.name} [{self.id}] end date is before rota start date: {self.end_date}")
+                Rota.add_warning(
+                    "Worker/Early end date",
+                    f"{self.name} [{self.id}] end date is before rota start date: {self.end_date}",
+                )
                 raise ValueError(f"End date must be after start date: {self.name}")
 
         # if no start date default to the start of the rota
@@ -132,7 +156,10 @@ class Worker(BaseModel):
                 self.calculated_start_date = Rota.start_date
             elif self.start_date >= Rota.rota_end_date:
                 # If it is after add a warning
-                Rota.add_warning("Worker/Late start date", f"{self.name} [{self.id}] start date is after rota end date: {self.start_date}")
+                Rota.add_warning(
+                    "Worker/Late start date",
+                    f"{self.name} [{self.id}] start date is after rota end date: {self.start_date}",
+                )
             else:
                 pass
 
@@ -152,18 +179,21 @@ class Worker(BaseModel):
 
             if date < self.calculated_start_date:
                 Rota.unavailable_to_work.add((self.id, week, day))
-                Rota.unavailable_to_work_reason[
+                Rota.unavailable_to_work_reason[(self.id, week, day)] = (
-                    (self.id, week, day)
+                    f"START DATE: {self.calculated_start_date}"
-                ] = f"START DATE: {self.calculated_start_date}"
+                )
 
             if date >= self.calculated_end_date:
                 Rota.unavailable_to_work.add((self.id, week, day))
-                Rota.unavailable_to_work_reason[
+                Rota.unavailable_to_work_reason[(self.id, week, day)] = (
-                    (self.id, week, day)
+                    f"END DATE: {self.calculated_end_date}"
-                ] = f"END DATE: {self.calculated_end_date}"
+                )
 
         if not self.not_available_to_work:
-            Rota.add_warning("Worker/No unavailability", f"{self.name} [{self.id}] has no unavailabilities (leave)")
+            Rota.add_warning(
+                "Worker/No unavailability",
+                f"{self.name} [{self.id}] has no unavailabilities (leave)",
+            )
 
         for item in self.oop:
             start_oop = item.start_date
@@ -207,9 +237,9 @@ class Worker(BaseModel):
                     ]:
                         week, day = weeks_days
                         Rota.unavailable_to_work.add((self.id, week, day))
-                        Rota.unavailable_to_work_reason[
+                        Rota.unavailable_to_work_reason[(self.id, week, day)] = (
-                            (self.id, week, day)
+                            f"OOP ({oop_name})".format(self.oop)
-                        ] = f"OOP ({oop_name})".format(self.oop)
+                        )
 
         days_to_end = (self.calculated_end_date - Rota.start_date).days
 
@@ -239,10 +269,11 @@ class Worker(BaseModel):
                 week = days_from_start // 7 + 1
                 day = Rota.days[(days_from_start % 7)]
                 Rota.unavailable_to_work.add((self.id, week, day))
-                Rota.unavailable_to_work_reason[(self.id, week, day)] = unavalability.reason
+                Rota.unavailable_to_work_reason[(self.id, week, day)] = (
+                    unavalability.reason
+                )
                 unavailable_set.add((week, day))
 
-
         for request in self.work_requests:
             days_from_start = (request.date - Rota.start_date).days
             week = days_from_start // 7 + 1
@@ -262,7 +293,9 @@ class Worker(BaseModel):
 
         if self.fte_adj > 100:
             # Shouldn't happen ? bug if it does
-            Rota.add_warning("FTE > 100%", f"{self.name} [{self.id}] FTE = {self.fte_adj}")
+            Rota.add_warning(
+                "FTE > 100%", f"{self.name} [{self.id}] FTE = {self.fte_adj}"
+            )
             raise ValueError("{} : fte_ajd = {}".format(self.name, self.fte_adj))
 
         # TODO: this has already been validated, consider moving
@@ -275,14 +308,22 @@ class Worker(BaseModel):
             if non_working_day.end_date is not None:
                 end_date = non_working_day.end_date
 
-            self.non_working_day_list.append((non_working_day.day, start_date, end_date))
+            self.non_working_day_list.append(
+                (non_working_day.day, start_date, end_date)
+            )
 
         if days_to_work < 1:
             self.fte_adj = 0
 
-        if days_to_work > 0 and (days_to_work - len(unavailable_set)) / days_to_work < 0.3:
+        if (
+            days_to_work > 0
+            and (days_to_work - len(unavailable_set)) / days_to_work < 0.3
+        ):
             console.print(f"{ self.name }: {days_to_work}, {Rota.rota_days_length}")
-            console.print(f"Warning {self.name} has excessive unavailability [{len(unavailable_set)}/{days_to_work}, adjusted fte={self.fte_adj}] (this may be infeasible)", style="red")
+            console.print(
+                f"Warning {self.name} has excessive unavailability [{len(unavailable_set)}/{days_to_work}, adjusted fte={self.fte_adj}] (this may be infeasible)",
+                style="red",
+            )
 
     def __lt__(self, other) -> bool:
         return (self.site, self.grade, self.fte_adj, self.name) < (
@@ -295,17 +336,17 @@ class Worker(BaseModel):
     def __str__(self) -> str:
         return f"{self.name}, {self.site}, {self.oop}"
 
-        #nwds = (
+        # nwds = (
         #    ", ".join([str(i) for i, start, end in self.non_working_day_list])
         #    if self.non_working_day_list is not None
         #    else ""
-        #)
+        # )
-        #return "{} {} [{}] REMOTE SITE:{}".format(
+        # return "{} {} [{}] REMOTE SITE:{}".format(
         #    self.name,
         #    (self.site, self.grade, self.fte_adj),
         #    nwds,
         #    self.remote_site,
-        #)
+        # )
 
     def get_details(self) -> str:
         return "{} {} {}".format(self.id, self.site[0], self.grade)

test/test_shifts.py

@@ -1,14 +1,13 @@
 import datetime
 import pytest
-from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, WarningTermination, days
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, WarningTermination, days, WorkerRequirement
 
-from rota.workers import NotAvailableToWork, WorkRequests, Worker
+from rota.workers import NotAvailableToWork, WorkRequests, Worker, generate_not_available_to_works
 
 import itertools
 
 
-def generate_basic_rota(weeks_to_rota=10, workers=2):
+def generate_basic_rota(weeks_to_rota=10, workers=2, start_date=datetime.date(2022, 3, 7)):
-    start_date = datetime.date(2022, 3, 7)
 
     Rota = RotaBuilder(
         start_date,
@@ -19,7 +18,6 @@ def generate_basic_rota(weeks_to_rota=10, workers=2):
     Rota.add_workers(
         [
             Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)
-            #Worker(name="worker2", site="group1", grade=1),
         ]
     )
 
@@ -547,56 +545,167 @@ class TestShiftDates:
         for worker in Rota.get_workers_by_group()["group2"]:
             assert Rota.get_worker_shifts_by_date(worker)[request_date] == "c"
 
-    def test_shift_start_date_end_date_block_paired(self):
+class TestShiftWorkerRequirements:
-        Rota = generate_basic_rota(workers=7)
+    def test_start_and_end_date(self):
+        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14))
 
+    def test_end_date_before_start_date(self):
+        with pytest.raises(ValueError):
+            WorkerRequirement(start_date=datetime.date(2022, 3, 14), end_date=datetime.date(2022, 3, 7))
 
-        request_date = Rota.start_date + datetime.timedelta(days=(7*6)+4)
+    def test_worker_requirement(self):
-
+        Rota = generate_basic_rota(workers=2)
-        no_work = [NotAvailableToWork(date=d, shift="*") for d in date_generator(request_date, 18)]
-
-        Rota.add_workers(
-            [
-                Worker(name=f"worker-wr-{i}", site="group2", grade=1, not_available_to_work=no_work) for i in range(1, 4)
-                #Worker(name="worker2", site="group1", grade=1),
-            ]
-        )
-        a, b, c = (SingleShift(
-            sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
-            workers_required=8,
-            constraint=[{"name": "max_shifts_per_week", "options": 4}],
-            end_date=(Rota.start_date + datetime.timedelta(days=7)),
-            #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-        ),
-        SingleShift(
-            sites=("group1", "group2"), name="b", length= 12.5, days=days[4:],
-            assign_as_block=True,
-            force_as_block=True,
-            workers_required=4,
-            constraint=[{"name": "max_shifts_per_week", "options": 4},
-                            {"name": "pre", "options": 3}, {"name": "post", "options": 3}],
-            start_date=(Rota.start_date + datetime.timedelta(days=7)),
-            end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-        ),
-        SingleShift(
-            sites=("group1", "group2"), name="c", length= 12.5, days=days[4:],
-            assign_as_block=True,
-            force_as_block=True,
-            workers_required=4,
-            constraint=[{"name": "max_shifts_per_week", "options": 4},
-                            {"name": "pre", "options": 3}, {"name": "post", "options": 3}],
-            start_date=(Rota.start_date + datetime.timedelta(days=(7*6))),
-        ))
-
 
+        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21))]
+        print(wr)
         Rota.add_shifts(
-            a, 
+            SingleShift(
-            b, c
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                assign_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))),
+            ),
         )
-        Rota.pair_shifts(b, c)
+        Rota.build_and_solve(options={"ratio": 0.1})
-
+        Rota.export_rota_to_html("test_worker_requirement")
-        Rota.build_and_solve(options={"ratio": 0.01})
-        Rota.export_rota_to_html("test_shift_start_date")
 
         assert Rota.results.solver.status == "ok"
-        assert Rota.results.solver.termination_condition == "optimal"
+
+        total_shifts = Rota.get_workers_total_shifts()
+
+        for worker in Rota.get_workers():
+            assert total_shifts[worker.name] == 4
+            assert "aaaa" in Rota.get_worker_shift_list_string(worker)
+
+
+    def test_worker_requirement2(self):
+        Rota = generate_basic_rota(workers=2)
+
+        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21)),
+                WorkerRequirement(start_date=datetime.date(2022, 3, 28))
+              ]
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                assign_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))),
+            ),
+        )
+        Rota.build_and_solve(options={"ratio": 0.1})
+        Rota.export_rota_to_html("test_worker_requirement")
+
+        assert Rota.results.solver.status == "ok"
+
+        total_shifts = Rota.get_workers_total_shifts()
+
+        for worker in Rota.get_workers():
+            assert total_shifts[worker.name] == 18
+            assert Rota.get_worker_shift_list_string(worker)[14:21] == "-------"
+
+    def test_worker_requirement_increase_workers(self):
+        Rota = generate_basic_rota(workers=2)
+
+        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
+                WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
+              ]
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                assign_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))),
+            ),
+        )
+        Rota.build_and_solve(options={"ratio": 0.1})
+        Rota.export_rota_to_html("test_worker_requirement")
+
+        assert Rota.results.solver.status == "ok"
+
+        total_shifts = Rota.get_workers_total_shifts()
+
+        for worker in Rota.get_workers():
+            assert total_shifts[worker.name] == 30
+            assert Rota.get_worker_shift_list_string(worker)[7:21] == "-" * 14
+            assert Rota.get_worker_shift_list_string(worker).endswith("aaaa---"*7)
+
+    def test_worker_requirement_balancing(self):
+        Rota = generate_basic_rota(workers=2)
+
+        natws = generate_not_available_to_works(datetime.date(2022,4,18), days=21, reason="holiday")
+
+        Rota.add_worker(
+            Worker(name=f"worker3", site="group1", grade=1, not_available_to_work=natws)
+        )
+
+        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
+                WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
+              ]
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                assign_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))),
+            ),
+        )
+        Rota.build_and_solve(options={"ratio": 0.0})
+        Rota.export_rota_to_html("test_worker_requirement")
+
+        assert Rota.results.solver.status == "ok"
+
+        total_shifts = Rota.get_workers_total_shifts()
+
+        for worker in Rota.get_workers():
+            assert total_shifts[worker.name] == 20
+            assert Rota.get_worker_shift_list_string(worker)[7:21] == "-" * 14
+
+            if worker.name == "worker3":
+                assert Rota.get_worker_shift_list_string(worker).endswith("aaaa---")
+                assert Rota.get_worker_shift_list_string(worker).startswith("aaaa---")
+
+
+    def test_worker_requirement_multiple_shifts(self):
+        Rota = generate_basic_rota(workers=4)
+
+        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
+                WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
+              ]
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[4:],
+                assign_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))),
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                assign_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))),
+            ),
+        )
+        Rota.build_and_solve(options={"ratio": 0.1})
+        Rota.export_rota_to_html("test_worker_requirement")
+
+        assert Rota.results.solver.status == "ok"
+
+        total_shifts = Rota.get_workers_total_shifts()
+
+        for worker in Rota.get_workers():
+            assert total_shifts[worker.name] == 30

test/test_workers.py

@@ -2,23 +2,29 @@ import pytest
 from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WarningTermination, days
 
 import datetime
-from rota.workers import Worker
+from rota.workers import NotAvailableToWork, Worker, generate_not_available_to_works
 
+def generate_basic_rota(weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7)):
 
-def setup_basic_rota():
+    Rota = RotaBuilder(
-    weeks_to_rota = 10
-    start_date = datetime.date(2022, 3, 7)
-
-    return RotaBuilder(
         start_date,
         weeks_to_rota=weeks_to_rota,
     )
 
+    # Add a few workers
+    if workers:
+        Rota.add_workers(
+            [
+                Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)
+            ]
+        )
+
+    return Rota
 
 class TestWorkers:
     def test_no_workers(self):
         """Check error raised if attempt to solve with no workers"""
-        Rota = setup_basic_rota()
+        Rota = generate_basic_rota()
 
         Rota.add_shifts(
             SingleShift(
@@ -43,7 +49,7 @@ class TestWorkers:
 
     def test_workers_duplicate_name(self):
 
-        Rota = setup_basic_rota()
+        Rota = generate_basic_rota()
 
         workers = []
         workers.append(Worker(name="worker1", site="group1", grade=1, fte=100))
@@ -325,7 +331,7 @@ class TestWorkers:
         Rota.export_rota_to_html("worker_oops")
 
     def test_worker_pairs(self):
-        Rota = setup_basic_rota()
+        Rota = generate_basic_rota()
 
         Rota.add_shifts(
             SingleShift(
@@ -395,4 +401,42 @@ class TestWorkers:
         Rota.export_rota_to_html("worker_pairs")
 
         assert Rota.results.solver.status == "ok"
-        assert Rota.get_workers_total_shifts()["worker3"] == 50
+        assert Rota.get_workers_total_shifts()["worker3"] == 50
+
+    def test_not_available_to_work(self):
+        
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",),
+                name="a",
+                length=12.5,
+                days=days,
+                workers_required=1,
+            ),
+        )
+
+        natws = generate_not_available_to_works(Rota.start_date, days=21)
+
+        natws.append(NotAvailableToWork(date=Rota.rota_end_date-datetime.timedelta(days=1)))
+
+
+
+        Rota.add_workers(
+            [
+                Worker(name="worker1", site="group1", grade=1),
+                Worker(name="worker2", site="group1", grade=1, not_available_to_work=natws),
+            ]
+        )
+
+        Rota.build_and_solve()
+        Rota.export_rota_to_html("test_not_available_to_work")
+
+        assert Rota.results.solver.status == "ok"
+
+        assert Rota.get_workers_total_shifts()["worker1"] == 35
+        assert Rota.get_workers_total_shifts()["worker2"] == 35
+
+        assert Rota.get_worker_shift_list_string(Rota.get_worker_by_name("worker2")).startswith("-" * 21)
+        assert Rota.get_worker_shift_list_string(Rota.get_worker_by_name("worker2")).endswith("-")