devplacepy/devplacepy/services/devii/tasks/schedule.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from typing import Literal, Optional

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

ISO_FORMAT = "%Y-%m-%dT%H:%M:%S"
CRON_MINUTES_LIMIT = 525600 * 4

ScheduleKind = Literal["once", "interval", "cron"]


def now_utc() -> datetime:
    return datetime.now(timezone.utc).replace(microsecond=0)


def to_iso(moment: datetime) -> str:
    if moment.tzinfo is None:
        moment = moment.replace(tzinfo=timezone.utc)
    return moment.astimezone(timezone.utc).strftime(ISO_FORMAT)


def from_iso(value: str) -> datetime:
    return datetime.strptime(value, ISO_FORMAT).replace(tzinfo=timezone.utc)


def _parse_field(field: str, low: int, high: int) -> set[int]:
    values: set[int] = set()
    for part in field.split(","):
        token, _, step_text = part.partition("/")
        step = int(step_text) if step_text else 1
        if step < 1:
            raise ValueError(f"Invalid step in cron field: {part}")
        if token == "*":
            start, end = low, high
        elif "-" in token:
            start_text, end_text = token.split("-", 1)
            start, end = int(start_text), int(end_text)
        else:
            start = end = int(token)
        if start < low or end > high or start > end:
            raise ValueError(f"Cron value out of range: {part}")
        values.update(range(start, end + 1, step))
    return values


def cron_next(expr: str, after: datetime) -> datetime:
    fields = expr.split()
    if len(fields) != 5:
        raise ValueError(
            "Cron expression must have 5 fields: minute hour dom month dow"
        )
    minutes = _parse_field(fields[0], 0, 59)
    hours = _parse_field(fields[1], 0, 23)
    days = _parse_field(fields[2], 1, 31)
    months = _parse_field(fields[3], 1, 12)
    weekdays_raw = _parse_field(fields[4], 0, 7)
    weekdays = {0 if value == 7 else value for value in weekdays_raw}
    day_restricted = fields[2] != "*"
    weekday_restricted = fields[4] != "*"

    candidate = after.replace(second=0, microsecond=0) + timedelta(minutes=1)
    for _ in range(CRON_MINUTES_LIMIT):
        if (
            candidate.minute in minutes
            and candidate.hour in hours
            and candidate.month in months
        ):
            cron_weekday = (candidate.weekday() + 1) % 7
            day_ok = candidate.day in days
            weekday_ok = cron_weekday in weekdays
            if day_restricted and weekday_restricted:
                match_day = day_ok or weekday_ok
            elif day_restricted:
                match_day = day_ok
            elif weekday_restricted:
                match_day = weekday_ok
            else:
                match_day = True
            if match_day:
                return candidate
        candidate += timedelta(minutes=1)
    raise ValueError("No cron match found within four years")


class Schedule(BaseModel):
    model_config = ConfigDict(extra="forbid")

    kind: ScheduleKind
    run_at: Optional[datetime] = None
    delay_seconds: Optional[int] = Field(default=None, ge=1)
    every_seconds: Optional[int] = Field(default=None, ge=1)
    start_at: Optional[datetime] = None
    cron: Optional[str] = None
    max_runs: Optional[int] = Field(default=None, ge=1)

    @field_validator("run_at", "start_at")
    @classmethod
    def _ensure_utc(cls, value: Optional[datetime]) -> Optional[datetime]:
        if value is None:
            return None
        if value.tzinfo is None:
            return value.replace(tzinfo=timezone.utc)
        return value.astimezone(timezone.utc)

    @model_validator(mode="after")
    def _validate_kind(self) -> "Schedule":
        if self.kind == "once" and self.run_at is None and self.delay_seconds is None:
            raise ValueError("kind=once requires run_at or delay_seconds")
        if self.kind == "interval" and self.every_seconds is None:
            raise ValueError("kind=interval requires every_seconds")
        if self.kind == "cron":
            if not self.cron:
                raise ValueError("kind=cron requires a cron expression")
            cron_next(self.cron, now_utc())
        return self

    def first_run(self, reference: datetime) -> datetime:
        if self.kind == "once":
            if self.run_at is not None:
                return self.run_at
            return reference + timedelta(seconds=self.delay_seconds or 0)
        if self.kind == "interval":
            if self.start_at is not None:
                return self.start_at
            return reference + timedelta(seconds=self.every_seconds or 0)
        return cron_next(self.cron or "", reference)

    def columns(self) -> dict[str, object]:
        return {
            "kind": self.kind,
            "run_at": to_iso(self.run_at) if self.run_at else None,
            "delay_seconds": self.delay_seconds,
            "every_seconds": self.every_seconds,
            "start_at": to_iso(self.start_at) if self.start_at else None,
            "cron": self.cron,
            "max_runs": self.max_runs,
        }


def next_run(
    kind: str, every_seconds: Optional[int], cron: Optional[str], reference: datetime
) -> Optional[datetime]:
    if kind == "interval" and every_seconds:
        return reference + timedelta(seconds=int(every_seconds))
    if kind == "cron" and cron:
        return cron_next(str(cron), reference)
    return None