scheduler/backend/test_window.py

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,
    daytime_shift2_minutes: int = 480,
) -> dict[int, int]:
    if daytime_shift2_only:
        daytime_minutes = max(0, min(int(daytime_shift2_minutes), 480))
        return {1: 0, 2: daytime_minutes, 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,
    daytime_shift2_minutes: int = 480,
) -> 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, daytime_shift2_minutes).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)