From e6bdd4a09e8762f331be90af9109c9cda6983f93 Mon Sep 17 00:00:00 2001 From: Ross Date: Mon, 5 Jan 2026 21:44:34 +0000 Subject: [PATCH] Add support for avoiding shifts on specific dates for workers - Introduced AvoidShiftOnDates model to define constraints for workers. - Enhanced Worker class to register avoid-shift rules with the rota builder. - Implemented logic in RotaBuilder to enforce these constraints during shift assignment. - Added tests to verify functionality for avoiding shifts on single dates and date ranges. --- leave.py | 6 ++ output/tests/timetable.css | 2 +- output/tests/timetable.js | 2 +- rota_generator/shifts.py | 122 ++++++++++++++++++++++++++++- rota_generator/workers.py | 40 ++++++++++ test/test_avoid_shifts_on_dates.py | 120 ++++++++++++++++++++++++++++ 6 files changed, 289 insertions(+), 3 deletions(-) create mode 100644 test/test_avoid_shifts_on_dates.py diff --git a/leave.py b/leave.py index 748da45..107e33a 100644 --- a/leave.py +++ b/leave.py @@ -259,3 +259,9 @@ def load_leave(Rota): # load_leave() + +def load_academy(Rota): + """ + Some trainees spend month blocks in the academy. This function loads that information from a google sheet that published it as a csv file. + """ + url = "https://docs.google.com/spreadsheets/d/e/2PACX-1vQtNrp9F1fjtAh8vFX-W_R3RlD9H-2P_vCLMoHrVKR0e24yabbsdrD65VjwxoHlV1Qnatmw1NeghQHG/pub?gid=1113837782&single=true&output=csv" \ No newline at end of file diff --git a/output/tests/timetable.css b/output/tests/timetable.css index 8a82cf4..6694d8b 120000 --- a/output/tests/timetable.css +++ b/output/tests/timetable.css @@ -1 +1 @@ -/home/ross/proc-rota/output/timetable.css \ No newline at end of file +/home/ross/proc_rota/output/timetable.css \ No newline at end of file diff --git a/output/tests/timetable.js b/output/tests/timetable.js index 7cdb8bf..b7adfdf 120000 --- a/output/tests/timetable.js +++ b/output/tests/timetable.js @@ -1 +1 @@ -/home/ross/proc-rota/output/timetable.js \ No newline at end of file +/home/ross/proc_rota/output/timetable.js \ No newline at end of file diff --git a/rota_generator/shifts.py b/rota_generator/shifts.py index 458438e..de8c3f8 100644 --- a/rota_generator/shifts.py +++ b/rota_generator/shifts.py @@ -300,10 +300,18 @@ class MaxShiftsPerWeekBlockConstraint(BaseShiftConstraint): max_shifts: int | None = None def _pydantic_to_dict(obj): + # Recursively convert pydantic models, lists and dicts to JSON-serializable structures + if obj is None: + return None if isinstance(obj, list): return [_pydantic_to_dict(x) for x in obj] + if isinstance(obj, dict): + return {k: _pydantic_to_dict(v) for k, v in obj.items()} + # datetime/date -> ISO string + if isinstance(obj, (datetime.date, datetime.datetime)): + return obj.isoformat() if hasattr(obj, "model_dump"): - return obj.model_dump() + return _pydantic_to_dict(obj.model_dump()) return obj @@ -582,6 +590,9 @@ class RotaConstraintOptions(BaseModel): # These are lists of dicts with keys like {"grades": [...], "weeks": [...], "shifts": [...]} avoid_shifts_by_grades: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning grades to shifts") avoid_shifts_by_worker_names: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning specific workers to shifts") + # Rules to avoid assigning specific workers to particular shifts on specific dates. + # Each entry is a dict with keys: {"names": [...], "shifts": [...], "start_date": date|None, "end_date": date|None, "dates": [date,...]|None} + avoid_shifts_by_worker_dates: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning specific workers to shifts on particular dates or date ranges") allocate_locum_shifts: bool = Field(True, description="Allow allocation of locum shifts") distribute_locum_shifts: bool = Field(True, description="Distribute locum shifts among available workers") balance_locum_shifts: bool = Field(False, description="Balance locum shifts similarly to normal shifts") @@ -2311,6 +2322,61 @@ class RotaBuilder(object): ) pass + # Date-based per-worker per-shift avoidance rules + for constraint_dict in self.constraint_options_model.avoid_shifts_by_worker_dates: + # validate shift names + if invalid_shifts := set(constraint_dict["shifts"]).difference(set(self.get_shift_names())): + message = f"avoid shifts by worker/dates constraint contains a non existent shift: {invalid_shifts}" + if self.ignore_valid_shifts: + self.add_warning("Non existent shift", message) + continue + else: + raise InvalidShift(message) + + if worker.name not in constraint_dict["names"]: + continue + + # Build set of allowed (week,day) pairs that fall within the date selector + selected_weeks_days = set() + + # explicit dates take precedence + if constraint_dict.get("dates"): + date_set = set(constraint_dict["dates"]) + for wk, d in self.get_week_day_combinations(): + date = self.get_date_by_week_day(wk, d) + if date in date_set: + selected_weeks_days.add((wk, d)) + else: + sd = constraint_dict.get("start_date") + ed = constraint_dict.get("end_date") + # If neither provided, apply to whole rota + if sd is None and ed is None: + for wk, d in self.get_week_day_combinations(): + selected_weeks_days.add((wk, d)) + else: + if sd is None: + sd = self.start_date + if ed is None: + ed = self.rota_end_date + for wk, d in self.get_week_day_combinations(): + date = self.get_date_by_week_day(wk, d) + if sd <= date <= ed: + selected_weeks_days.add((wk, d)) + + if not selected_weeks_days: + # nothing to constrain + continue + + # Add constraint that forbids the given shifts for this worker on the selected dates + self.model.constraints.add( + 0 + == sum( + self.model.works[worker.id, wk, d, shift] + for wk, d, shift in self.get_all_shiftname_combinations() + if (wk, d) in selected_weeks_days and shift in constraint_dict["shifts"] + ) + ) + for constraint_dict in self.constraint_options_model.avoid_shifts_by_worker_names: if invalid_shifts := set(constraint_dict["shifts"]).difference( set(self.get_shift_names()) @@ -4092,6 +4158,44 @@ class RotaBuilder(object): {"names": names, "weeks": weeks, "shifts": shifts} ) + def add_avoid_shifts_for_workers_on_dates( + self, + names: Iterable[str], + shifts: Iterable[str], + start_date: datetime.date | str | None = None, + end_date: datetime.date | str | None = None, + dates: Iterable[datetime.date] | None = None, + ): + """Add a rule preventing the given workers from being assigned the given shifts + on the provided date range or explicit dates. + + - `names`: iterable of worker names to which the rule applies. + - `shifts`: iterable of shift names to avoid. + - `start_date`/`end_date`: inclusive date range (accepts date or string "%d/%m/%Y" or "%d/%m/%y"). + - `dates`: optional explicit iterable of `datetime.date` objects. + """ + def _coerce_date(d): + if d is None: + return None + if isinstance(d, datetime.date): + return d + if isinstance(d, str): + for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"): + try: + return datetime.datetime.strptime(d, fmt).date() + except ValueError: + continue + raise ValueError(f"Cannot parse date: {d}") + + sd = _coerce_date(start_date) + ed = _coerce_date(end_date) + + entry = {"names": list(names), "shifts": list(shifts), "start_date": sd, "end_date": ed, "dates": None} + if dates is not None: + entry["dates"] = [d if isinstance(d, datetime.date) else _coerce_date(d) for d in dates] + + self.constraint_options_model.avoid_shifts_by_worker_dates.append(entry) + def add_grade_constraint_by_week( self, grades: Iterable[int], weeks: Iterable[int], shifts ): @@ -4900,6 +5004,22 @@ class RotaBuilder(object): timetable[worker.name][week][day] = assigned_shifts return timetable + def get_assigned_shifts_for_worker_on_day(self, worker_id: int, week: int, day: DayStr) -> list: + """Return a list of assigned SingleShift objects for the given worker on a specific week/day. + + This is a small compatibility helper used by tests that expect objects rather than shift-name strings. + """ + works = self.model.works + assigned = [] + for shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True): + try: + if works[worker_id, week, day, shift_name].value > 0.5: + assigned.append(self.get_shift_by_name(shift_name)) + except Exception: + # If variable missing or access error, skip + continue + return assigned + def get_worker_timetable_brief( self, show_prefs=False, marker_every=30, show_unavailable=False ): diff --git a/rota_generator/workers.py b/rota_generator/workers.py index dface81..fbf915c 100644 --- a/rota_generator/workers.py +++ b/rota_generator/workers.py @@ -111,6 +111,29 @@ class OutOfProgramme(BaseModel): reason: str = "" +class AvoidShiftOnDates(BaseModel): + """Worker-level constraint: avoid listed shifts on given dates or date range.""" + shifts: list[str] + start_date: Optional[datetime.date] = None + end_date: Optional[datetime.date] = None + dates: Optional[list[datetime.date]] = None + + @field_validator("start_date", "end_date", mode="before") + def coerce_dates(cls, v): + # allow strings in common formats + if v is None: + return None + if isinstance(v, datetime.date): + return v + if isinstance(v, str): + for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"): + try: + return datetime.datetime.strptime(v, fmt).date() + except Exception: + continue + raise ValueError(f"Cannot parse date: {v}") + + class Worker(BaseModel): name: str site: str @@ -154,6 +177,7 @@ class Worker(BaseModel): shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist weekend_shift_target_number: int = 0 + avoid_shifts_on_dates: list[AvoidShiftOnDates] = [] @@ -395,6 +419,22 @@ class Worker(BaseModel): style="red", ) + # Register any per-worker avoid-shift-on-dates rules with the rota builder + try: + for av in getattr(self, "avoid_shifts_on_dates", []): + entry = { + "names": [self.name], + "shifts": av.shifts, + "start_date": av.start_date, + "end_date": av.end_date, + "dates": av.dates, + } + # append to rota-level list; the model builder will enforce these + if hasattr(Rota.constraint_options_model, "avoid_shifts_by_worker_dates"): + Rota.constraint_options_model.avoid_shifts_by_worker_dates.append(entry) + except Exception: + Rota.add_warning("Worker/avoid_shifts_on_dates", f"Failed to register avoid_shifts_on_dates for {self.name}") + def __lt__(self, other) -> bool: return (self.site, self.grade, self.fte_adj, self.name) < ( other.site, diff --git a/test/test_avoid_shifts_on_dates.py b/test/test_avoid_shifts_on_dates.py new file mode 100644 index 0000000..88991c4 --- /dev/null +++ b/test/test_avoid_shifts_on_dates.py @@ -0,0 +1,120 @@ +import datetime +from rota_generator.shifts import RotaBuilder, SingleShift, days +from rota_generator.workers import Worker, AvoidShiftOnDates + + +def build_simple_rota(): + start = datetime.date(2026, 4, 6) # Monday + Rota = RotaBuilder(start, weeks_to_rota=4, name="test_rota") + Rota.add_shifts( + SingleShift(sites=("exeter",), name="night_weekday", length=12.25, days=days[:4], workers_required=1), + SingleShift(sites=("exeter",), name="plymouth_twilight", length=8, days=days[5:], workers_required=1), + ) + Rota.build_shifts() + return Rota + + +def test_avoid_single_date_for_worker(): + Rota = build_simple_rota() + # create a worker who should avoid night_weekday on 2026-04-07 + avoid = AvoidShiftOnDates(shifts=["night_weekday"], dates=[datetime.date(2026, 4, 7)]) + w = Worker(name="Test Worker", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid]) + Rota.add_worker(w) + Rota.build_workers() + + # find week/day corresponding to 2026-04-07 + weeks_days = [(wk, d) for wk, d in Rota.weeks_days_product if Rota.week_day_date_map[(wk, d)] == datetime.date(2026, 4, 7)] + assert len(weeks_days) == 1 + wk, day = weeks_days[0] + + # The model should mark works[worker.id, wk, day, 'night_weekday'] impossible -> check unavailable_to_work + assert (w.id, wk, day) not in Rota.unavailable_to_work + + # Build and solve minimal model (but we can inspect constraints). Instead, check that the avoid rule was registered + found = False + for c in Rota.constraint_options_model.avoid_shifts_by_worker_dates: + if "names" in c and w.name in c["names"]: + found = True + # ensure the date is represented + assert c.get("dates") is not None and datetime.date(2026,4,7) in c.get("dates") + assert found + + +def test_avoid_date_range_for_worker_applies_to_multiple_days(): + Rota = build_simple_rota() + # Avoid night_weekday for 2026-04-06 through 2026-04-09 + avoid = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="09/04/2026") + w = Worker(name="Range Worker", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid]) + Rota.add_worker(w) + Rota.build_workers() + + # Collect dates between start and end + sd = datetime.date(2026,4,6) + ed = datetime.date(2026,4,9) + dates_in_range = [d for d in Rota.get_date_range(sd, ed)] + assert len(dates_in_range) == 4 + + # Ensure the rota-level constraint was added + found = False + for c in Rota.constraint_options_model.avoid_shifts_by_worker_dates: + if w.name in c.get("names", []): + found = True + assert c.get("start_date") == sd + assert c.get("end_date") == ed + assert found + + +def test_export_rota_for_visual_inspection(): + Rota = build_simple_rota() + # add two workers + avoid = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="26/04/2026") + w1 = Worker(name="Export A", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid]) + w2 = Worker(name="Export B", site="exeter", grade=3, fte=100) + + + + Rota.add_workers((w1, w2)) + + # Build model but do not solve, then export + Rota.name = "export_test_rota" + Rota.build_and_solve(solve=True) + + assert Rota.results.solver.status == "ok", "Rota should be feasible" + + # Next we check that worker w1 has no night_weekday shifts assigned between 06/04/2026 and 26/04/2026 + sd = datetime.date(2026,4,6) + ed = datetime.date(2026,4,26) + for single_date in Rota.get_date_range(sd, ed): + wk_day = Rota.get_week_day_by_date(single_date) + assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w1.id, wk_day[0], wk_day[1]) + for sh in assigned_shifts: + assert sh.name != "night_weekday", f"Worker {w1.name} was assigned forbidden shift on {single_date}" + + # We also check that w1 has 4 night shifts in total (the max possible in 4 weeks) + night_shift_count = 0 + for wk, day in Rota.weeks_days_product: + assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w1.id, wk, day) + for sh in assigned_shifts: + if sh.name == "night_weekday": + night_shift_count += 1 + assert night_shift_count == 4, f"Worker {w1.name} should have 4 night shifts, has {night_shift_count}" + + # And that w2 has 12 night shifts (no restrictions) + night_shift_count_w2 = 0 + for wk, day in Rota.weeks_days_product: + assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w2.id, wk, day) + for sh in assigned_shifts: + if sh.name == "night_weekday": + night_shift_count_w2 += 1 + assert night_shift_count_w2 == 12, f"Worker {w2.name} should have 12 night shifts, has {night_shift_count_w2}" + + # Next check that adding a short avoid to w2 results in an infeasible model + avoid2 = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="07/04/2026") + w2.avoid_shifts_on_dates.append(avoid2) + Rota.build_and_solve(solve=True) + + assert Rota.results.solver.termination_condition == "infeasible", "Rota should be infeasible due to avoid shifts on dates constraints" + + + Rota.export_rota_to_html("avoid_shifts_on_dates_test.html", folder="tests") +