|
| 1 | +from collections.abc import Mapping |
| 2 | +from dataclasses import dataclass, field |
| 3 | +from enum import Enum |
| 4 | +from typing import Annotated, Any |
| 5 | + |
| 6 | +from pydantic import PlainSerializer, PlainValidator, WithJsonSchema |
| 7 | + |
| 8 | + |
| 9 | +class CronError(ValueError): |
| 10 | + """Raised on invalid cron expression.""" |
| 11 | + |
| 12 | + |
| 13 | +class WeekDays(str, Enum): |
| 14 | + """Weekday names.""" |
| 15 | + |
| 16 | + SUN = "SUN" |
| 17 | + MON = "MON" |
| 18 | + TUE = "TUE" |
| 19 | + WED = "WED" |
| 20 | + THU = "THU" |
| 21 | + FRI = "FRI" |
| 22 | + SAT = "SAT" |
| 23 | + |
| 24 | + |
| 25 | +class Months(str, Enum): |
| 26 | + """Month names.""" |
| 27 | + |
| 28 | + JAN = "JAN" |
| 29 | + FEB = "FEB" |
| 30 | + MAR = "MAR" |
| 31 | + APR = "APR" |
| 32 | + MAY = "MAY" |
| 33 | + JUN = "JUN" |
| 34 | + JUL = "JUL" |
| 35 | + AUG = "AUG" |
| 36 | + SEP = "SEP" |
| 37 | + OCT = "OCT" |
| 38 | + NOV = "NOV" |
| 39 | + DEC = "DEC" |
| 40 | + |
| 41 | + |
| 42 | +@dataclass(frozen=True) |
| 43 | +class FieldSpec: |
| 44 | + """Allowed values of a cron field.""" |
| 45 | + |
| 46 | + name: str |
| 47 | + low: int |
| 48 | + high: int |
| 49 | + aliases: Mapping[str, int] = field(default_factory=dict) |
| 50 | + |
| 51 | + |
| 52 | +MONTH_NAMES: dict[str, int] = {n.value: i for i, n in enumerate(Months, start=1)} |
| 53 | +WEEKDAY_NAMES: dict[str, int] = {n.value: i for i, n in enumerate(WeekDays)} |
| 54 | + |
| 55 | +MINUTE = FieldSpec("minute", 0, 59) |
| 56 | +HOUR = FieldSpec("hour", 0, 23) |
| 57 | +DAY = FieldSpec("day", 1, 31) |
| 58 | +MONTH = FieldSpec("month", 1, 12, MONTH_NAMES) |
| 59 | +WEEKDAY = FieldSpec("weekday", 0, 7, WEEKDAY_NAMES) # 0 and 7 = SUN |
| 60 | + |
| 61 | + |
| 62 | +def _is_uint(s: str) -> bool: |
| 63 | + return s.isascii() and s.isdigit() |
| 64 | + |
| 65 | + |
| 66 | +@dataclass(frozen=True) |
| 67 | +class CronField: |
| 68 | + """ |
| 69 | + Single cron field. |
| 70 | +
|
| 71 | + Accepted values: |
| 72 | + *: any value |
| 73 | + N: one value |
| 74 | + A-B: a range, A <= B |
| 75 | + */S: a step within the whole range |
| 76 | + A-B/S: a step inside the range |
| 77 | + Comma separated lists of the above. |
| 78 | + """ |
| 79 | + |
| 80 | + value: str |
| 81 | + spec: FieldSpec |
| 82 | + |
| 83 | + def __post_init__(self) -> None: |
| 84 | + for part in self.value.split(","): |
| 85 | + self._check_part(part) |
| 86 | + |
| 87 | + def _check_part(self, part: str) -> None: |
| 88 | + name = self.spec.name |
| 89 | + |
| 90 | + base, has_step, step = part.partition("/") |
| 91 | + if has_step: |
| 92 | + if not _is_uint(step) or int(step) == 0: |
| 93 | + raise CronError(f"{name}: the step should be > 0, got {step!r}") |
| 94 | + if base != "*" and "-" not in base: |
| 95 | + raise CronError( |
| 96 | + f"{name}: the step is allowed only after '*' or a range, " |
| 97 | + f"got {part!r}", |
| 98 | + ) |
| 99 | + |
| 100 | + if base == "*": |
| 101 | + return |
| 102 | + |
| 103 | + start, has_range, end = base.partition("-") |
| 104 | + lo = self._to_number(start) |
| 105 | + if has_range: |
| 106 | + hi = self._to_number(end) |
| 107 | + if lo > hi: |
| 108 | + raise CronError(f"{name}: start must be lower than the end {part!r}") |
| 109 | + |
| 110 | + def _to_number(self, token: str) -> int: |
| 111 | + spec = self.spec |
| 112 | + alias = token.upper() |
| 113 | + if alias in spec.aliases: |
| 114 | + return spec.aliases[alias] |
| 115 | + if not _is_uint(token): |
| 116 | + raise CronError(f"{spec.name}: a number expected, got {token!r}") |
| 117 | + n = int(token) |
| 118 | + if not spec.low <= n <= spec.high: |
| 119 | + raise CronError(f"{spec.name}: {n} is out of range {spec.low}-{spec.high}") |
| 120 | + return n |
| 121 | + |
| 122 | + def to_numeric(self) -> str: |
| 123 | + """Convert the field to numbers only, in the form pycron understands.""" |
| 124 | + is_weekday = self.spec is WEEKDAY |
| 125 | + parts: list[str] = [] |
| 126 | + for part in self.value.split(","): |
| 127 | + base, _, step = part.partition("/") |
| 128 | + if base == "*": |
| 129 | + parts.append(part) |
| 130 | + continue |
| 131 | + start, has_range, end = base.partition("-") |
| 132 | + lo = self._to_number(start) |
| 133 | + if not has_range: |
| 134 | + parts.append(str(lo % 7 if is_weekday else lo)) |
| 135 | + elif is_weekday: |
| 136 | + values = range(lo, self._to_number(end) + 1, int(step or 1)) |
| 137 | + parts.extend(str(v % 7) for v in values) |
| 138 | + else: |
| 139 | + suffix = f"/{step}" if step else "" |
| 140 | + parts.append(f"{lo}-{self._to_number(end)}{suffix}") |
| 141 | + return ",".join(parts) |
| 142 | + |
| 143 | + |
| 144 | +@dataclass |
| 145 | +class Cron: |
| 146 | + """Validated 5-field cron expression.""" |
| 147 | + |
| 148 | + minute: CronField |
| 149 | + hour: CronField |
| 150 | + day: CronField |
| 151 | + month: CronField |
| 152 | + weekday: CronField |
| 153 | + |
| 154 | + def __post_init__(self) -> None: |
| 155 | + expected = (MINUTE, HOUR, DAY, MONTH, WEEKDAY) |
| 156 | + actual = (self.minute, self.hour, self.day, self.month, self.weekday) |
| 157 | + for f, spec in zip(actual, expected, strict=True): |
| 158 | + if f.spec is not spec: |
| 159 | + raise CronError( |
| 160 | + f"field {spec.name} is created with wrong spec: {f.spec.name}", |
| 161 | + ) |
| 162 | + |
| 163 | + @classmethod |
| 164 | + def parse(cls, expr: str) -> "Cron": |
| 165 | + """ |
| 166 | + Parse cron string. |
| 167 | +
|
| 168 | + :param expr: cron expression with 5 fields. |
| 169 | + :raises CronError: if the expression is invalid. |
| 170 | + :return: cron object. |
| 171 | + """ |
| 172 | + parts = expr.split() |
| 173 | + if len(parts) != 5: |
| 174 | + raise CronError(f"Need 5 fields, got: {len(parts)}: {expr!r}") |
| 175 | + |
| 176 | + specs = (MINUTE, HOUR, DAY, MONTH, WEEKDAY) |
| 177 | + return cls(*(CronField(v, s) for v, s in zip(parts, specs, strict=True))) |
| 178 | + |
| 179 | + def to_pycron(self) -> str: |
| 180 | + """Cron string with numeric values only, suitable for pycron.""" |
| 181 | + fields = (self.minute, self.hour, self.day, self.month, self.weekday) |
| 182 | + return " ".join(f.to_numeric() for f in fields) |
| 183 | + |
| 184 | + def __str__(self) -> str: |
| 185 | + return ( |
| 186 | + f"{self.minute.value} {self.hour.value} {self.day.value} " |
| 187 | + f"{self.month.value} {self.weekday.value}" |
| 188 | + ) |
| 189 | + |
| 190 | + |
| 191 | +def _validate_cron(value: Any) -> Cron: |
| 192 | + if isinstance(value, Cron): |
| 193 | + return value |
| 194 | + if isinstance(value, str): |
| 195 | + return Cron.parse(value) |
| 196 | + raise CronError(f"cron must be a string or Cron, got {type(value).__name__}") |
| 197 | + |
| 198 | + |
| 199 | +CronExpression = Annotated[ |
| 200 | + Cron | str, |
| 201 | + PlainValidator(_validate_cron), |
| 202 | + PlainSerializer(str, return_type=str), |
| 203 | + WithJsonSchema({"type": "string"}), |
| 204 | +] |
0 commit comments