diff --git a/requirements.txt b/requirements.txt index b3af670..05217d0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,4 +13,5 @@ python-dotenv requests loguru numpy -highspy \ No newline at end of file +highspy +typer diff --git a/rota/shifts.py b/rota/shifts.py index b9f5c13..a96c3db 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, 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: - extra = Extra.allow - orm_mode = True + model_config = ConfigDict( + extra = "allow", + #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,18 +1726,22 @@ class RotaBuilder(object): for constraint_shift in self.get_shifts_with_constraints( "max_shifts_per_week" ): - self.model.constraints.add( - constraint_shift.constraint_options["max_shifts_per_week"] - >= sum( - 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( - constraint_shift + try: + self.model.constraints.add( + constraint_shift.constraint_options["max_shifts_per_week"] + >= sum( + 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( + constraint_shift + ) + if w == week ) - 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: - print(s.name, self.week_day_date_map[(week, day)]) + 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) @@ -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,13 +2834,25 @@ 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) """ - return self.shifts + 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): shifts = [shift for shift in self.shifts if worker_site in shift.sites] @@ -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 = [] diff --git a/rota/workers.py b/rota/workers.py index bd9b431..2032f16 100644 --- a/rota/workers.py +++ b/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: - extra = Extra.allow + model_config = ConfigDict( + extra = "allow", + ) # def __init__( # self, diff --git a/test/test_shifts.py b/test/test_shifts.py index 3ff1ff8..9f9a0df 100644 --- a/test/test_shifts.py +++ b/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="worker2", 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), ] ) @@ -311,4 +311,131 @@ class TestShiftConstraints: Rota.build_and_solve(options={"ratio": 0.000}) Rota.export_rota_to_html("post") - assert Rota.results.solver.status in ("warning", "error") \ No newline at end of file + 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) \ No newline at end of file