ross/proc-rota
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add worker requirements

Browse Source
This commit is contained in:
Ross
2025-01-01 21:58:48 +00:00
parent 9af01a3415
commit 7ce4677921
5 changed files with 535 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rota/shifts.py
+171 -64
View File
@@ -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())