Compare commits

..
2 Commits
Author SHA1 Message Date
Ross d66ddf29ad Add support for avoiding shifts on specific dates with reasons 2026-01-05 22:07:23 +00:00
Ross e6bdd4a09e 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.
2026-01-05 21:44:34 +00:00
8 changed files with 487 additions and 7 deletions
+4 -2
View File
@@ -49,7 +49,7 @@ def main(
time_to_run: int = 60 * 60 * 4,
ratio: float = 0.001,
start_date: datetime.datetime = "2026-03-02",
weeks: int = 12,
weeks: int = 26,
bom: int = 2,
):
rota_start_date = start_date.date()
@@ -276,7 +276,7 @@ def main(
)
# Prevent shifts for ST2s for 2 weeks
Rota.add_grade_constraint_by_week([2], [1, 2], [shift.name for shift in Rota.shifts])
#Rota.add_grade_constraint_by_week([2], [1, 2], [shift.name for shift in Rota.shifts])
Rota.add_min_summed_grade_by_shifts_per_day_constraint(["weekend_plymouth1", "weekend_plymouth2"], 6)
Rota.add_min_summed_grade_by_shifts_per_day_constraint(["plymouth_twilight", "plymouth_bank_holidays"], 6)
@@ -463,10 +463,12 @@ def main(
print(w)
Rota.add_worker(w)
leave.load_academy(Rota)
# Rota.build_workers()
# Rota.build_model()
solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
# solver_options = {"seconds": time_to_run, "threads": 10}
+95 -1
View File
@@ -1,5 +1,5 @@
import csv
from datetime import datetime
from datetime import datetime, timedelta
import re
import sys
from requests import Session
@@ -259,3 +259,97 @@ 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"
with Session() as s:
headers = {"Cache-Control": "no-cache", "Pragma": "no-cache"}
download = s.get(url, headers=headers)
decoded = download.content.decode("utf-8")
reader = csv.reader(decoded.splitlines(), delimiter=',')
n = 0
names = []
# collect date rows (these are the month start boundaries) and their cells
date_rows: list[datetime.date] = []
date_row_cells: list[list[str]] = []
for row in reader:
n += 1
if n == 1:
# header row: first cell is column title, remaining are worker names
names = [c.strip() for c in row[1:]]
continue
if not row:
continue
row_title = row[0].strip()
cells = row[1:]
# If this row is a date row like '02/03/2026' then treat it as a month-start boundary
if re.match(date_re, row_title):
try:
try:
d = datetime.strptime(row_title, "%d/%m/%Y").date()
except ValueError:
d = datetime.strptime(row_title, "%d/%m/%y").date()
except Exception:
continue
date_rows.append(d)
# normalize cells length to match names
row_cells = [ (cells[i] if i < len(cells) else "").strip() for i in range(len(names)) ]
date_row_cells.append(row_cells)
# If no date rows found, nothing to do
if not date_rows:
return
# Identify twilight shift names available in the rota
twilight_shift_names = [s.name for s in Rota.get_shifts() if "twilight" in s.name]
# For each date boundary, compute its end (next boundary - 1 day) and register avoids where cell=='academy'
for idx, start_date in enumerate(date_rows):
if idx + 1 < len(date_rows):
end_date = date_rows[idx + 1] - timedelta(days=1)
else:
end_date = Rota.rota_end_date
cells = date_row_cells[idx]
for i, cell in enumerate(cells):
if i >= len(names):
break
v = (cell or "").lower()
if "academy" in v:
worker_name = names[i]
matching_workers = [w for w in Rota.workers if w.name.strip() == worker_name.strip()]
if not matching_workers:
logger.warning(f"Academy data: worker '{worker_name}' not found in rota, skipping")
continue
worker = matching_workers[0]
# choose twilight shifts relevant to the worker's site
candidate_shifts = []
for sh_name in twilight_shift_names:
try:
sh = Rota.get_shift_by_name(sh_name)
except Exception:
continue
if hasattr(worker, "site") and worker.site in sh.sites:
candidate_shifts.append(sh_name)
elif sh_name.startswith(f"{worker.site}_"):
candidate_shifts.append(sh_name)
if not candidate_shifts:
candidate_shifts = twilight_shift_names.copy()
# register as a date range with a reason 'academy'
Rota.add_avoid_shifts_for_workers_on_dates(
names=[worker.name], shifts=candidate_shifts, start_date=start_date, end_date=end_date, reason="academy"
)
+1 -1
View File
@@ -1 +1 @@
/home/ross/proc-rota/output/timetable.css
/home/ross/proc_rota/output/timetable.css
+1 -1
View File
@@ -1 +1 @@
/home/ross/proc-rota/output/timetable.js
/home/ross/proc_rota/output/timetable.js
+48
View File
@@ -259,6 +259,52 @@ function generateExtra() {
$(table).before(summary_button);
// Add per-table avoid-toggle button (near the summary button)
avoid_button = $("<button class='auto-generated'>Toggle Avoids</button>").click((evt) => {
// ensure style exists once
if ($('#avoid-shift-style').length === 0) {
$('head').append(`
<style id='avoid-shift-style'>
td.rota-day.avoid-shift.avoid-highlight { border: 2px solid rgba(255,140,0,0.9) !important; box-shadow: 0 0 6px rgba(255,140,0,0.25); }
</style>
`);
}
let $btn = $(evt.target);
let enabled = $btn.data('enabled') || false;
enabled = !enabled;
$btn.data('enabled', enabled);
let $table = $(table);
$table.find("td.rota-day.avoid-shift").each((i, el) => {
let $el = $(el);
if (enabled) {
$el.addClass('avoid-highlight');
// preserve original title
if (!$el.data('orig-title')) $el.data('orig-title', $el.attr('title') || '');
let avoid_shifts = $el.attr('data-avoid-shifts') || '';
let avoid_reason = $el.attr('data-avoid-reason') || '';
let suffix = [];
// make the constraint type explicit in the tooltip
suffix.push(`constraint: avoid_shifts_by_worker_dates`);
if (avoid_shifts) suffix.push(`shifts: ${avoid_shifts}`);
if (avoid_reason) suffix.push(`reason: ${avoid_reason}`);
if (suffix.length) {
let newTitle = $el.data('orig-title') + ' [' + suffix.join('; ') + ']';
$el.attr('title', newTitle);
}
} else {
$el.removeClass('avoid-highlight');
// restore original title
if ($el.data('orig-title') !== undefined) $el.attr('title', $el.data('orig-title'));
}
});
$btn.text(enabled ? 'Hide Avoids' : 'Toggle Avoids');
});
$(table).before(avoid_button);
});
@@ -281,6 +327,8 @@ $("#toggle-prior-adjust").on("click", (evt) => {
$("table.tsummary").find(`td.prior-adjust, th.prior-adjust-col`).toggle();
});
// Note: per-table avoid-toggle buttons are added next to each table's summary button inside generateExtra().
$("table.tsummary").append("<tr class='header'></tr>");
h = $("table.tsummary tr.header");
h.append(`<th>Worker</th>`);
+176 -2
View File
@@ -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,45 @@ 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,
reason: str | 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, "reason": reason}
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 +5005,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
):
@@ -4984,6 +5105,49 @@ class RotaBuilder(object):
):
model = self.model
# Build a lookup of per-worker/day avoid-shift rules (worker.id, week, day) -> {shifts: [...], reason: str}
avoid_lookup: dict[tuple, dict] = {}
try:
for entry in getattr(self.constraint_options_model, "avoid_shifts_by_worker_dates", []):
names = entry.get("names", []) or []
shifts = entry.get("shifts", []) or []
reason = entry.get("reason")
dates = entry.get("dates")
sd = entry.get("start_date")
ed = entry.get("end_date")
date_list = []
if dates:
date_list = dates
elif sd is not None or ed is not None:
try:
for d in self.get_date_range(sd, ed):
date_list.append(d)
except Exception:
# fall back to trying to iterate
pass
for name in names:
matching_workers = [w for w in self.workers if w.name.strip() == name.strip()]
if not matching_workers:
continue
wid = matching_workers[0].id
for d in date_list:
try:
wk, day = self.get_week_day_by_date(d)
except Exception:
continue
key = (wid, wk, day)
existing = avoid_lookup.get(key, {"shifts": [], "reason": None})
existing["shifts"].extend(shifts)
# preserve reason if provided
if reason:
existing["reason"] = reason
avoid_lookup[key] = existing
except Exception:
# If constraint options not present or parsing fails, continue without avoid info
avoid_lookup = {}
timetable = []
if self.run_start_time is not None and self.run_end_time is not None:
@@ -5132,6 +5296,16 @@ class RotaBuilder(object):
if locum:
css_class = " ".join((css_class, "locum-shift"))
# Add visual marker/data if this day/worker has avoid-shift rules
avoid_info = avoid_lookup.get((worker.id, week, day))
avoid_attr = ""
if avoid_info is not None:
# mark the cell as having avoids; include which shifts and the reason
avoid_shifts = ",".join(avoid_info.get("shifts", []))
avoid_reason = avoid_info.get("reason") or ""
css_class = " ".join((css_class, "avoid-shift"))
avoid_attr = f" data-avoid-shifts='{avoid_shifts}' data-avoid-reason='{avoid_reason}'"
remote_site = ""
if shift_name:
shift = self.get_shift_by_name(shift_name)
@@ -5147,7 +5321,7 @@ class RotaBuilder(object):
f" data-available='{available}'"
f" data-unavailable_reason='{unavailable_reason}'"
f" data-date='{d}' data-week='{week}' data-day='{day}'"
f"{remote_site}{requests}{bank_holiday}>{shift_display_char}</td>"
f"{remote_site}{requests}{bank_holiday}{avoid_attr}>{shift_display_char}</td>"
)
shift_count = ""
+42
View File
@@ -111,6 +111,30 @@ 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
reason: Optional[str] = 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 +178,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 +420,23 @@ 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,
"reason": getattr(av, "reason", None),
}
# 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,
+120
View File
@@ -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")