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.
This commit is contained in:
+121
-1
@@ -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
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user