.
This commit is contained in:
Ross
@@ -0,0 +1,475 @@
|
||||
import datetime
|
||||
from collections import defaultdict
|
||||
from typing import Iterable, List, Literal, Optional, Set
|
||||
from pydantic import BaseModel, ConfigDict, field_validator, Field
|
||||
|
||||
from rich.pretty import pprint
|
||||
from rota_generator.console import console
|
||||
|
||||
# from .shifts import RotaBuilder, days, sites
|
||||
import uuid
|
||||
|
||||
# from rota.shifts import RotaBuilder
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from rota_generator.shifts import RotaBuilder
|
||||
|
||||
days = Literal["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
|
||||
whole_days = [
|
||||
"sunday",
|
||||
"monday",
|
||||
"tuesday",
|
||||
"wednesday",
|
||||
"thursday",
|
||||
"friday",
|
||||
"saturday",
|
||||
]
|
||||
|
||||
|
||||
class NotAvailableToWork(BaseModel):
|
||||
date: datetime.date
|
||||
reason: str = "unknown"
|
||||
|
||||
class HardDayDependency(BaseModel):
|
||||
days: Set[str] = Field(..., description="List of days for the dependency group")
|
||||
weeks: Optional[List[int]] = None # If None, applies to all weeks
|
||||
start_date: Optional[datetime.date] = None # Optional start date for the dependency
|
||||
end_date: Optional[datetime.date] = None # Optional end date for the dependency
|
||||
|
||||
@field_validator("weeks")
|
||||
def weeks_none_if_dates_set(cls, weeks, values):
|
||||
data = values.data
|
||||
if (data.get("start_date") is not None or data.get("end_date") is not None) and weeks is not None:
|
||||
raise ValueError("If start_date or end_date is set, weeks must not be set")
|
||||
return weeks
|
||||
|
||||
class HardDayExclusion(BaseModel):
|
||||
days: set[str] = Field(..., description="List of days to exclude together")
|
||||
weeks: Optional[List[int]] = None # If None, applies to all weeks
|
||||
start_date: Optional[datetime.date] = None # Optional start date for the dependency
|
||||
end_date: Optional[datetime.date] = None # Optional end date for the dependency
|
||||
|
||||
@field_validator("weeks")
|
||||
def weeks_none_if_dates_set(cls, weeks, values):
|
||||
data = values.data
|
||||
if (data.get("start_date") is not None or data.get("end_date") is not None) and weeks is not None:
|
||||
raise ValueError("If start_date or end_date is set, weeks must not be set")
|
||||
return weeks
|
||||
|
||||
def generate_not_available_to_works(
|
||||
start_date: datetime.date,
|
||||
end_date: datetime.date | None = None,
|
||||
days: int | None = None,
|
||||
reason: str = "unknown",
|
||||
) -> List[NotAvailableToWork]:
|
||||
if end_date is None and days is None:
|
||||
raise ValueError("Either end_date or days must be provided")
|
||||
if end_date is not None and days is not None:
|
||||
raise ValueError("Only one of end_date or days must be provided")
|
||||
|
||||
if end_date is not None:
|
||||
days = (end_date - start_date).days
|
||||
|
||||
not_available_to_works = []
|
||||
for day in range(days):
|
||||
not_available_to_works.append(NotAvailableToWork(date=start_date + datetime.timedelta(days=day), reason=reason))
|
||||
|
||||
return not_available_to_works
|
||||
|
||||
|
||||
class NonWorkingDays(BaseModel):
|
||||
day: days
|
||||
start_date: datetime.date | None = None
|
||||
end_date: datetime.date | None = None
|
||||
|
||||
@field_validator("day")
|
||||
def day_in(cls, day):
|
||||
for whole_day in whole_days:
|
||||
if day.lower() in whole_day:
|
||||
return whole_day[:3].capitalize()
|
||||
raise ValueError("Invalid day")
|
||||
|
||||
|
||||
class WorkRequests(BaseModel):
|
||||
"""
|
||||
Shift is the name of the shift requested, * is a wildcard for any shift
|
||||
"""
|
||||
date: datetime.date
|
||||
shift: str = ""
|
||||
|
||||
|
||||
class PreferenceNotToWork(BaseModel):
|
||||
date: datetime.date
|
||||
reason: str = ""
|
||||
|
||||
|
||||
class OutOfProgramme(BaseModel):
|
||||
start_date: datetime.date
|
||||
end_date: datetime.date
|
||||
reason: str = ""
|
||||
|
||||
|
||||
class Worker(BaseModel):
|
||||
name: str
|
||||
site: str
|
||||
# Grade are equivalent to roles, by keeping them integer defining model
|
||||
# rules is easier.
|
||||
grade: int = 1
|
||||
# We can either have a user generated ID
|
||||
id: Optional[int | uuid.UUID | str] = None
|
||||
fte: int = 100
|
||||
nwds: list[NonWorkingDays] = []
|
||||
start_date: datetime.date | None = None
|
||||
end_date: datetime.date | None = None
|
||||
oop: list[OutOfProgramme] = []
|
||||
not_available_to_work: list[NotAvailableToWork] = []
|
||||
pref_not_to_work: list[PreferenceNotToWork] = []
|
||||
work_requests: list[WorkRequests] = []
|
||||
locum_availability: list[WorkRequests] = []
|
||||
locum_max_shifts: int = 0
|
||||
locum_max_shifts_per_week: int = 0
|
||||
locum_on_nwds: bool = False
|
||||
remote_site: str = "plymouth" # We set a default proc_site
|
||||
previous_shifts: dict = {}
|
||||
shift_balance_extra: dict = {}
|
||||
bank_holiday_extra: int = 0
|
||||
pair: int | str | None = None
|
||||
locum: bool = False
|
||||
allowed_multi_shift_sets: list[set] = [] # or list/set of frozensets
|
||||
prefer_multi_shift_together: int = 0 # Default: no preference
|
||||
# Set to a positive integer to prefer, negative to discourage, 0 for neutral
|
||||
hard_day_dependencies: list[HardDayDependency] = []
|
||||
hard_day_exclusions: list[HardDayExclusion] = []
|
||||
force_assign_with: dict[str, list[str]] = {} # e.g. {"oncall": ["weekend"]}
|
||||
max_shifts_per_week_by_shift_name: dict[str, int] = {} # e.g. {"night": 2, "a": 3}
|
||||
forced_assignments: List[tuple[int, str, str]] = [] # (week, day, shift_name)
|
||||
forced_assignments_by_date: List[tuple[datetime.date, str]] = [] # (date, shift_name)
|
||||
max_days_worked_per_week: int | None= None # Default: no limit, can be set to a specific number
|
||||
|
||||
|
||||
assign_as_block_preferences: dict[str, float] = {} # shift_name -> weight (positive = prefer block)
|
||||
|
||||
shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist
|
||||
|
||||
weekend_shift_target_number: int = 0
|
||||
|
||||
|
||||
|
||||
model_config = ConfigDict(
|
||||
extra="allow",
|
||||
)
|
||||
|
||||
# def __init__(
|
||||
# self,
|
||||
# #Rota: "RotaBuilder", # Aim to remove (workers should be independent of rotas)
|
||||
# ):
|
||||
#
|
||||
# self.name = name
|
||||
# self.site = site
|
||||
# self.grade = grade
|
||||
# self.fte = fte
|
||||
# self.nwd = nwd
|
||||
# self.pair = pair
|
||||
# self.shift_balance_extra = shift_balance_extra
|
||||
# self.bank_holiday_extra = bank_holiday_extra
|
||||
# self.start_date
|
||||
#
|
||||
# self.remote_site = remote_site
|
||||
#
|
||||
# self.previous_shifts = previous_shifts
|
||||
#
|
||||
#
|
||||
# # days_to_work = Rota.rota_days_length
|
||||
|
||||
def load_rota(self, Rota: "RotaBuilder"):
|
||||
self.proportion_rota_to_work: float = 1
|
||||
self.shift_target_number: dict[str, int] = defaultdict(int)
|
||||
|
||||
if self.id is None:
|
||||
self.id = uuid.uuid4()
|
||||
|
||||
if self.start_date is not None and self.end_date is not None:
|
||||
if self.start_date >= self.end_date:
|
||||
Rota.add_warning(
|
||||
"Worker/Early end date",
|
||||
f"{self.name} [{self.id}] end date is before rota start date: {self.end_date}",
|
||||
)
|
||||
raise ValueError(f"End date must be after start date: {self.name}")
|
||||
|
||||
# if no start date default to the start of the rota
|
||||
if self.start_date is None:
|
||||
self.calculated_start_date = Rota.start_date
|
||||
else:
|
||||
self.calculated_start_date = self.start_date
|
||||
# ? test if start date is valid
|
||||
if self.start_date < Rota.start_date:
|
||||
# If worker start date is before rota start date use the rota start date
|
||||
self.calculated_start_date = Rota.start_date
|
||||
elif self.start_date >= Rota.rota_end_date:
|
||||
# If it is after add a warning
|
||||
Rota.add_warning(
|
||||
"Worker/Late start date",
|
||||
f"{self.name} [{self.id}] start date is after rota end date: {self.start_date}",
|
||||
)
|
||||
else:
|
||||
pass
|
||||
|
||||
if self.end_date is None:
|
||||
self.calculated_end_date = Rota.rota_end_date
|
||||
else:
|
||||
self.calculated_end_date = self.end_date
|
||||
# self.calculated_end_date = datetime.datetime.strptime(end_date, "%d/%m/%y").date()
|
||||
|
||||
if self.calculated_end_date > Rota.rota_end_date:
|
||||
self.calculated_end_date = Rota.rota_end_date
|
||||
|
||||
days_to_work = (self.calculated_end_date - self.calculated_start_date).days
|
||||
|
||||
for week, day in Rota.weeks_days_product:
|
||||
date = Rota.week_day_date_map[(week, day)]
|
||||
|
||||
if date < self.calculated_start_date:
|
||||
Rota.unavailable_to_work.add((self.id, week, day))
|
||||
Rota.unavailable_to_work_reason[(self.id, week, day)] = (
|
||||
f"START DATE: {self.calculated_start_date}"
|
||||
)
|
||||
|
||||
if date >= self.calculated_end_date:
|
||||
Rota.unavailable_to_work.add((self.id, week, day))
|
||||
Rota.unavailable_to_work_reason[(self.id, week, day)] = (
|
||||
f"END DATE: {self.calculated_end_date}"
|
||||
)
|
||||
|
||||
if not self.not_available_to_work:
|
||||
Rota.add_warning(
|
||||
"Worker/No unavailability",
|
||||
f"{self.name} [{self.id}] has no unavailabilities (leave)",
|
||||
)
|
||||
|
||||
for item in self.oop:
|
||||
start_oop = item.start_date
|
||||
end_oop = item.end_date
|
||||
oop_name = item.reason
|
||||
# start_oop, end_oop = oop
|
||||
if isinstance(start_oop, datetime.date):
|
||||
start_oop_date = start_oop
|
||||
else:
|
||||
start_oop_date = datetime.datetime.strptime(
|
||||
start_oop, "%d/%m/%y"
|
||||
).date()
|
||||
|
||||
if isinstance(end_oop, datetime.date):
|
||||
end_oop_date = end_oop
|
||||
else:
|
||||
end_oop_date = datetime.datetime.strptime(end_oop, "%d/%m/%y").date()
|
||||
|
||||
if start_oop_date >= end_oop_date:
|
||||
raise ValueError(f"End OOP date must be after start date [{self.name} - Start date: {self.start_date} / End date: {self.end_date}]")
|
||||
|
||||
# ignore oops if they finish before the rota (or worker) start date
|
||||
if end_oop_date > self.calculated_start_date:
|
||||
if start_oop_date > self.calculated_end_date:
|
||||
pass
|
||||
else:
|
||||
if end_oop_date > Rota.rota_end_date:
|
||||
end_oop_date = Rota.rota_end_date
|
||||
|
||||
if start_oop_date < Rota.start_date:
|
||||
start_oop_date = Rota.start_date
|
||||
|
||||
oop_length = (end_oop_date - start_oop_date).days
|
||||
print(f"oop length {oop_length}")
|
||||
days_to_work = days_to_work - oop_length
|
||||
|
||||
days_until_oop = (start_oop_date - Rota.start_date).days
|
||||
|
||||
for weeks_days in Rota.weeks_days_product[
|
||||
days_until_oop : days_until_oop + oop_length
|
||||
]:
|
||||
week, day = weeks_days
|
||||
Rota.unavailable_to_work.add((self.id, week, day))
|
||||
Rota.unavailable_to_work_reason[(self.id, week, day)] = (
|
||||
f"OOP ({oop_name})".format(self.oop)
|
||||
)
|
||||
|
||||
days_to_end = (self.calculated_end_date - Rota.start_date).days
|
||||
|
||||
# loop throught dates converting to week / day combination
|
||||
for preference in self.pref_not_to_work:
|
||||
days_from_start = (preference.date - Rota.start_date).days
|
||||
# Ignore dates past the end of the rota (or end date)
|
||||
if days_from_start < days_to_end:
|
||||
week = days_from_start // 7 + 1
|
||||
day = Rota.days[(days_from_start % 7)]
|
||||
# Weight the value to take into account the number of preferences
|
||||
# 1 is added to the total number of requests to ensure they do not outweight
|
||||
# a single (or fewer) request(s)
|
||||
Rota.pref_not_to_work[(self.id, week, day)] = 1 / (
|
||||
len(self.pref_not_to_work) + 1
|
||||
)
|
||||
Rota.pref_not_to_work_reason[(self.id, week, day)] = preference.reason
|
||||
|
||||
# print(not_available_to_work)
|
||||
# loop throught dates converting to week / day combination
|
||||
|
||||
unavailable_set = set()
|
||||
for unavalability in self.not_available_to_work:
|
||||
days_from_start = (unavalability.date - Rota.start_date).days
|
||||
# Ignore dates past the end of the rota (or end date)
|
||||
if days_from_start < days_to_end:
|
||||
week = days_from_start // 7 + 1
|
||||
day = Rota.days[(days_from_start % 7)]
|
||||
Rota.unavailable_to_work.add((self.id, week, day))
|
||||
Rota.unavailable_to_work_reason[(self.id, week, day)] = (
|
||||
unavalability.reason
|
||||
)
|
||||
unavailable_set.add((week, day))
|
||||
|
||||
for request in self.work_requests:
|
||||
days_from_start = (request.date - Rota.start_date).days
|
||||
week = days_from_start // 7 + 1
|
||||
day = Rota.days[(days_from_start % 7)]
|
||||
if request.shift == "*":
|
||||
for shift in Rota.get_shifts_for_worker_site(self.site):
|
||||
Rota.work_requests.add((self.id, week, day, shift.name))
|
||||
else:
|
||||
Rota.work_requests.add((self.id, week, day, request.shift))
|
||||
|
||||
for request in self.locum_availability:
|
||||
days_from_start = (request.date - Rota.start_date).days
|
||||
week = days_from_start // 7 + 1
|
||||
day = Rota.days[(days_from_start % 7)]
|
||||
if request.shift == "*":
|
||||
# We may want to add a limit to shifts that can be filled by locums?
|
||||
for shift in Rota.get_shifts():
|
||||
Rota.locum_availability.add((self.id, week, day, shift.name))
|
||||
else:
|
||||
Rota.locum_availability.add((self.id, week, day, request.shift))
|
||||
|
||||
# Calculate the proportion of the rota that is being worked
|
||||
self.proportion_rota_to_work = days_to_work / Rota.rota_days_length
|
||||
self.days_to_work = days_to_work
|
||||
|
||||
# We have to adjust the full time equivalent for people who CCT / leave the rota early
|
||||
self.fte_adj = self.fte * self.proportion_rota_to_work
|
||||
|
||||
self.fte_adj_shifts = {}
|
||||
if self.shift_fte_overrides:
|
||||
for shift, fte in self.shift_fte_overrides.items():
|
||||
self.fte_adj_shifts[shift] = fte * self.proportion_rota_to_work
|
||||
|
||||
|
||||
if self.fte_adj > 100:
|
||||
# Shouldn't happen ? bug if it does
|
||||
Rota.add_warning(
|
||||
"FTE > 100%", f"{self.name} [{self.id}] FTE = {self.fte_adj}"
|
||||
)
|
||||
raise ValueError("{} : fte_ajd = {}".format(self.name, self.fte_adj))
|
||||
|
||||
# TODO: this has already been validated, consider moving
|
||||
self.non_working_day_list = []
|
||||
for non_working_day in self.nwds:
|
||||
start_date = Rota.start_date
|
||||
end_date = Rota.rota_end_date
|
||||
if non_working_day.start_date is not None:
|
||||
start_date = non_working_day.start_date
|
||||
if non_working_day.end_date is not None:
|
||||
end_date = non_working_day.end_date
|
||||
|
||||
self.non_working_day_list.append(
|
||||
(non_working_day.day, start_date, end_date)
|
||||
)
|
||||
|
||||
if days_to_work < 1:
|
||||
self.fte_adj = 0
|
||||
|
||||
if (
|
||||
days_to_work > 0
|
||||
and (days_to_work - len(unavailable_set)) / days_to_work < 0.3
|
||||
):
|
||||
console.print(f"{ self.name }: {days_to_work}, {Rota.rota_days_length}")
|
||||
console.print(
|
||||
f"Warning {self.name} has excessive unavailability [{len(unavailable_set)}/{days_to_work}, adjusted fte={self.fte_adj}] (this may be infeasible)",
|
||||
style="red",
|
||||
)
|
||||
|
||||
def __lt__(self, other) -> bool:
|
||||
return (self.site, self.grade, self.fte_adj, self.name) < (
|
||||
other.site,
|
||||
other.grade,
|
||||
other.fte_adj,
|
||||
other.name,
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.name}, {self.site}, {self.oop}"
|
||||
|
||||
# nwds = (
|
||||
# ", ".join([str(i) for i, start, end in self.non_working_day_list])
|
||||
# if self.non_working_day_list is not None
|
||||
# else ""
|
||||
# )
|
||||
# return "{} {} [{}] REMOTE SITE:{}".format(
|
||||
# self.name,
|
||||
# (self.site, self.grade, self.fte_adj),
|
||||
# nwds,
|
||||
# self.remote_site,
|
||||
# )
|
||||
|
||||
def get_details(self) -> str:
|
||||
return "{} {} {}".format(self.id, self.site[0], self.grade)
|
||||
|
||||
def get_full_details(self) -> str:
|
||||
return "{}/{}/{}".format(self.site, self.grade, self.name)
|
||||
|
||||
def get_shift_targets(self) -> dict:
|
||||
return self.shift_target_number
|
||||
|
||||
def get_fte(self, shift: str | None=None) -> int:
|
||||
if shift is not None:
|
||||
if shift in self.fte_adj_shifts:
|
||||
return self.fte_adj_shifts[shift]
|
||||
return self.fte_adj
|
||||
|
||||
def allow_shifts_together(self, *shifts):
|
||||
if not hasattr(self, "allowed_multi_shift_sets"):
|
||||
self.allowed_multi_shift_sets = []
|
||||
self.allowed_multi_shift_sets.append(frozenset(shifts))
|
||||
|
||||
def add_hard_day_dependency(self, days: set[str], weeks: Optional[List[int]] = None, start_date: Optional[datetime.date] = None, end_date: Optional[datetime.date] = None):
|
||||
"""
|
||||
Add a hard day dependency for this worker.
|
||||
days: List of days that must be grouped together.
|
||||
weeks: Optional list of weeks this dependency applies to. If None, applies to all weeks.
|
||||
"""
|
||||
self.hard_day_dependencies.append(HardDayDependency(days=days, weeks=weeks, start_date=start_date, end_date=end_date))
|
||||
|
||||
def add_hard_day_exclusion(self, days: set[str], weeks: Optional[List[int]] = None, start_date: Optional[datetime.date] = None, end_date: Optional[datetime.date] = None):
|
||||
"""
|
||||
Add a hard day exclusion for this worker.
|
||||
days: List of days that must be excluded together.
|
||||
weeks: Optional list of weeks this exclusion applies to. If None, applies to all weeks.
|
||||
"""
|
||||
self.hard_day_exclusions.append(HardDayExclusion(days=days, weeks=weeks, start_date=start_date, end_date=end_date))
|
||||
|
||||
def add_force_assign_with(self, shift: str, with_shifts: list[str] | str):
|
||||
if not hasattr(self, "force_assign_with"):
|
||||
self.force_assign_with = {}
|
||||
if isinstance(with_shifts, str):
|
||||
with_shifts = [with_shifts]
|
||||
self.force_assign_with[shift] = with_shifts
|
||||
|
||||
def set_max_shifts_per_week(self, shift_name: str, max_per_week: int):
|
||||
"""Helper to set the maximum number of shifts of a given type per week for this worker."""
|
||||
if not hasattr(self, "max_shifts_per_week_by_shift_name"):
|
||||
self.max_shifts_per_week_by_shift_name = {}
|
||||
self.max_shifts_per_week_by_shift_name[shift_name] = max_per_week
|
||||
|
||||
def force_assign_shift(self, week: int, day: str, shift_name: str):
|
||||
"""Force this worker to be assigned to a shift on a specific week/day."""
|
||||
self.forced_assignments.append((week, day, shift_name))
|
||||
|
||||
def force_assign_shift_by_date(self, date: datetime.date, shift_name: str):
|
||||
"""Force this worker to be assigned to a shift on a specific date."""
|
||||
self.forced_assignments_by_date.append((date, shift_name))
|
||||
Reference in New Issue
Block a user