diff --git a/gen.py b/gen.py index 2ecd894..8b5c116 100644 --- a/gen.py +++ b/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, - #force_as_block=True, + # rota_on_nwds=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(), - ), - 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(), + #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(), + # ), ) # 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_individually(solver_options) + # Rota.solve_shifts_by_block(solver_options, block_length=13) + # 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") diff --git a/rota/shifts.py b/rota/shifts.py index 7aa7e68..0d3fa24 100644 --- a/rota/shifts.py +++ b/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", - #orm_mode = True + extra="allow", + # 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( - sum( - self.model.blocks_worker_shift_assigned[ - worker.id, week, shift.name - ] - for worker in self.workers - ) - <= shift.workers_required - ) + if workers_required: + self.model.constraints.add( + sum( + self.model.blocks_worker_shift_assigned[ + worker.id, week, shift.name + ] + for worker in self.workers + ) + <= 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 - self.model.constraints.add( - self.model.shift_week_worker_assigned[ - shift.name, week, worker.id - ] - * len(shift.days) - == sum( - self.model.works[worker.id, week, day, shift.name] - for day in shift.days - ) - ) + if shift.get_worker_requirement_by_date( + self.get_week_start_date(week) + ): + # Force shifts to be assigned in blocks + try: + self.model.constraints.add( + self.model.shift_week_worker_assigned[ + shift.name, week, worker.id + ] + * 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) - #try: - #if constraint_shift.start_date <= week_start_date < constraint_shift.end_date: + # week_start_date = self.get_week_start_date(week) + # try: + # 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) - #if constraint_shift.start_date <= week_start_date < constraint_shift.end_date: - #works = 1 - #else: + # week_start_date = self.get_week_start_date(week) + # if constraint_shift.start_date <= week_start_date < constraint_shift.end_date: + # works = 1 + # 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()) diff --git a/rota/workers.py b/rota/workers.py index 2032f16..85bd1f3 100644 --- a/rota/workers.py +++ b/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[ - (self.id, week, day) - ] = f"START DATE: {self.calculated_start_date}" + Rota.unavailable_to_work_reason[(self.id, week, day)] = ( + 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[ - (self.id, week, day) - ] = f"END DATE: {self.calculated_end_date}" + Rota.unavailable_to_work_reason[(self.id, week, day)] = ( + 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[ - (self.id, week, day) - ] = f"OOP ({oop_name})".format(self.oop) + Rota.unavailable_to_work_reason[(self.id, week, day)] = ( + 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) diff --git a/test/test_shifts.py b/test/test_shifts.py index 8e13f92..e306f7d 100644 --- a/test/test_shifts.py +++ b/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): - start_date = datetime.date(2022, 3, 7) +def generate_basic_rota(weeks_to_rota=10, workers=2, 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): - Rota = generate_basic_rota(workers=7) +class TestShiftWorkerRequirements: + 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) - - 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))), - )) - + def test_worker_requirement(self): + Rota = generate_basic_rota(workers=2) + wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21))] + print(wr) Rota.add_shifts( - a, - b, c + 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.pair_shifts(b, c) - - Rota.build_and_solve(options={"ratio": 0.01}) - Rota.export_rota_to_html("test_shift_start_date") + Rota.build_and_solve(options={"ratio": 0.1}) + Rota.export_rota_to_html("test_worker_requirement") assert Rota.results.solver.status == "ok" - assert Rota.results.solver.termination_condition == "optimal" \ No newline at end of file + + 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 \ No newline at end of file diff --git a/test/test_workers.py b/test/test_workers.py index ca4be0d..50bbb39 100644 --- a/test/test_workers.py +++ b/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(): - weeks_to_rota = 10 - start_date = datetime.date(2022, 3, 7) - - return RotaBuilder( + Rota = 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 \ No newline at end of file + 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("-") \ No newline at end of file