68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from datetime import date, timedelta
|
||
|
|
|
||
|
|
|
||
|
|
def is_off_day(current_date: date, holiday_dates: set[str]) -> bool:
|
||
|
|
return current_date.weekday() >= 5 or current_date.isoformat() in holiday_dates
|
||
|
|
|
||
|
|
def get_shift_sequence(
|
||
|
|
current_date: date,
|
||
|
|
holiday_dates: set[str],
|
||
|
|
daytime_shift2_only: bool = False,
|
||
|
|
) -> list[tuple[date, int]]:
|
||
|
|
is_holiday = current_date.isoformat() in holiday_dates
|
||
|
|
weekday = current_date.weekday()
|
||
|
|
shifts: list[tuple[date, int]] = []
|
||
|
|
|
||
|
|
if daytime_shift2_only and weekday < 5 and not is_holiday:
|
||
|
|
return [(current_date, 2)]
|
||
|
|
|
||
|
|
if is_off_day(current_date, holiday_dates):
|
||
|
|
return [(current_date, 1), (current_date, 2), (current_date, 3)]
|
||
|
|
|
||
|
|
shifts.append((current_date, 3))
|
||
|
|
next_day = current_date + timedelta(days=1)
|
||
|
|
if is_off_day(next_day, holiday_dates):
|
||
|
|
cursor = next_day
|
||
|
|
while is_off_day(cursor, holiday_dates):
|
||
|
|
shifts.extend([(cursor, 1), (cursor, 2), (cursor, 3)])
|
||
|
|
cursor += timedelta(days=1)
|
||
|
|
shifts.append((cursor, 1))
|
||
|
|
return shifts
|
||
|
|
|
||
|
|
shifts.append((next_day, 1))
|
||
|
|
return shifts
|
||
|
|
|
||
|
|
|
||
|
|
def get_shift_capacity_for_date(
|
||
|
|
current_date: date,
|
||
|
|
holiday_dates: set[str],
|
||
|
|
daytime_shift2_only: bool = False,
|
||
|
|
) -> dict[int, int]:
|
||
|
|
if daytime_shift2_only:
|
||
|
|
return {1: 0, 2: 480, 3: 0}
|
||
|
|
if is_off_day(current_date, holiday_dates):
|
||
|
|
return {1: 480, 2: 480, 3: 480}
|
||
|
|
return {1: 480, 2: 0, 3: 480}
|
||
|
|
|
||
|
|
|
||
|
|
def get_shift_sequence_with_capacity(
|
||
|
|
current_date: date,
|
||
|
|
holiday_dates: set[str],
|
||
|
|
daytime_shift2_only: bool = False,
|
||
|
|
) -> tuple[list[tuple[date, int]], int]:
|
||
|
|
shifts = get_shift_sequence(current_date, holiday_dates, daytime_shift2_only)
|
||
|
|
capacity = sum(
|
||
|
|
get_shift_capacity_for_date(day, holiday_dates, daytime_shift2_only).get(shift_index, 0)
|
||
|
|
for day, shift_index in shifts
|
||
|
|
)
|
||
|
|
return shifts, capacity
|
||
|
|
|
||
|
|
|
||
|
|
def next_window_start_date(shifts: list[tuple[date, int]]) -> date:
|
||
|
|
last_date, last_shift = shifts[-1]
|
||
|
|
if last_shift in (1, 2):
|
||
|
|
return last_date
|
||
|
|
return last_date + timedelta(days=1)
|