1from __future__ import annotations
2
3import datetime
4import subprocess
5from itertools import chain
6
7from plain.utils import timezone
8
9from .jobs import Job
10from .registry import jobs_registry, register_job
11
12__all__ = ["Schedule", "ScheduledCommand"]
13
14_MONTH_NAMES = {
15 "JAN": 1,
16 "FEB": 2,
17 "MAR": 3,
18 "APR": 4,
19 "MAY": 5,
20 "JUN": 6,
21 "JUL": 7,
22 "AUG": 8,
23 "SEP": 9,
24 "OCT": 10,
25 "NOV": 11,
26 "DEC": 12,
27}
28_DAY_NAMES = {
29 "SUN": 0,
30 "MON": 1,
31 "TUE": 2,
32 "WED": 3,
33 "THU": 4,
34 "FRI": 5,
35 "SAT": 6,
36}
37
38
39class _ScheduleComponent:
40 def __init__(self, values: list[int], raw: str | int = "") -> None:
41 self.values = sorted(values)
42 self._raw = raw
43
44 def __str__(self) -> str:
45 if self._raw:
46 return str(self._raw)
47 return ",".join(str(v) for v in self.values)
48
49 def __eq__(self, other: object) -> bool:
50 return isinstance(other, _ScheduleComponent) and self.values == other.values
51
52 @property
53 def is_wildcard(self) -> bool:
54 """Whether this field contains a ``*``.
55
56 Cron's day-of-month/day-of-week OR rule treats a field as restricted
57 only when it has no ``*``, so a stepped wildcard like ``*/2`` counts as
58 unrestricted here too.
59 """
60 return "*" in str(self._raw)
61
62 @classmethod
63 def parse(
64 cls,
65 value: int | str,
66 min_allowed: int,
67 max_allowed: int,
68 str_conversions: dict[str, int] | None = None,
69 ) -> _ScheduleComponent:
70 if str_conversions is None:
71 str_conversions = {}
72
73 if isinstance(value, int):
74 if value < min_allowed or value > max_allowed:
75 raise ValueError(
76 f"Schedule component should be between {min_allowed} and {max_allowed}"
77 )
78 return cls([value], raw=value)
79
80 if not isinstance(value, str):
81 raise TypeError("Schedule component should be an int or str")
82
83 # First split any subcomponents and re-parse them
84 if "," in value:
85 return cls(
86 list(
87 chain.from_iterable(
88 cls.parse(
89 sub_value, min_allowed, max_allowed, str_conversions
90 ).values
91 for sub_value in value.split(",")
92 )
93 ),
94 raw=value,
95 )
96
97 if value == "*":
98 return cls(list(range(min_allowed, max_allowed + 1)), raw=value)
99
100 def _convert(value: str) -> int:
101 result = str_conversions.get(value.upper(), value)
102 return int(result)
103
104 if "/" in value:
105 values, step = value.split("/")
106 values = cls.parse(values, min_allowed, max_allowed, str_conversions)
107 return cls([v for v in values.values if v % int(step) == 0], raw=value)
108
109 if "-" in value:
110 start, end = value.split("-")
111 return cls(list(range(_convert(start), _convert(end) + 1)), raw=value)
112
113 return cls([_convert(value)], raw=value)
114
115
116class Schedule:
117 def __init__(
118 self,
119 *,
120 minute: int | str = "*",
121 hour: int | str = "*",
122 day_of_month: int | str = "*",
123 month: int | str = "*",
124 day_of_week: int | str = "*",
125 combine_days_with_or: bool = False,
126 raw: str = "",
127 ) -> None:
128 self.minute = _ScheduleComponent.parse(minute, min_allowed=0, max_allowed=59)
129 self.hour = _ScheduleComponent.parse(hour, min_allowed=0, max_allowed=23)
130 self.day_of_month = _ScheduleComponent.parse(
131 day_of_month, min_allowed=1, max_allowed=31
132 )
133 self.month = _ScheduleComponent.parse(
134 month,
135 min_allowed=1,
136 max_allowed=12,
137 str_conversions=_MONTH_NAMES,
138 )
139 # Cron numbers weekdays Sunday=0..Saturday=6 and also accepts 7 as an
140 # alias for Sunday. Parse over 0..7, then fold 7 into 0 so the values
141 # line up with the cron weekday computed in next().
142 parsed_days = _ScheduleComponent.parse(
143 day_of_week,
144 min_allowed=0,
145 max_allowed=7,
146 str_conversions=_DAY_NAMES,
147 )
148 self.day_of_week = _ScheduleComponent(
149 sorted({0 if value == 7 else value for value in parsed_days.values}),
150 raw=day_of_week,
151 )
152
153 # Standard cron runs a job when *either* the day-of-month or the
154 # day-of-week matches, but only when both fields are restricted. That
155 # quirk is faithful to cron strings; the keyword API keeps the more
156 # obvious AND semantics unless you opt in here.
157 self.combine_days_with_or = combine_days_with_or
158
159 self._raw = raw
160
161 def __str__(self) -> str:
162 if self._raw:
163 return self._raw
164 return f"{self.minute} {self.hour} {self.day_of_month} {self.month} {self.day_of_week}"
165
166 def __repr__(self) -> str:
167 return f"<Schedule {self}>"
168
169 @classmethod
170 def from_cron(cls, cron: str) -> Schedule:
171 raw = cron
172
173 if cron == "@yearly" or cron == "@annually":
174 cron = "0 0 1 1 *"
175 elif cron == "@monthly":
176 cron = "0 0 1 * *"
177 elif cron == "@weekly":
178 cron = "0 0 * * 0"
179 elif cron == "@daily" or cron == "@midnight":
180 cron = "0 0 * * *"
181 elif cron == "@hourly":
182 cron = "0 * * * *"
183
184 minute, hour, day_of_month, month, day_of_week = cron.split()
185
186 return cls(
187 minute=minute,
188 hour=hour,
189 day_of_month=day_of_month,
190 month=month,
191 day_of_week=day_of_week,
192 combine_days_with_or=True,
193 raw=raw,
194 )
195
196 def next(self, now: datetime.datetime | None = None) -> datetime.datetime:
197 """
198 Find the next datetime that matches the schedule after the given datetime.
199 """
200 dt = now or timezone.localtime() # Use the defined plain timezone by default
201
202 # We only care about minutes, so immediately jump to the next minute
203 dt += datetime.timedelta(minutes=1)
204 dt = dt.replace(second=0, microsecond=0)
205
206 def _go_to_next_day(v: datetime.datetime) -> datetime.datetime:
207 v = v + datetime.timedelta(days=1)
208 return v.replace(
209 hour=self.hour.values[0],
210 minute=self.minute.values[0],
211 )
212
213 # If we don't find a value in the next 500 days,
214 # then the schedule is probably never going to match (i.e. Feb 31)
215 max_future = dt + datetime.timedelta(days=500)
216
217 # Only combine the two day fields with OR when both are restricted.
218 # This doesn't depend on the candidate day, so decide it once up front.
219 use_or_for_days = (
220 self.combine_days_with_or
221 and not self.day_of_month.is_wildcard
222 and not self.day_of_week.is_wildcard
223 )
224
225 while True:
226 # Cron numbers weekdays Sunday=0..Saturday=6; Python's weekday() is
227 # Monday=0..Sunday=6, so shift it before comparing.
228 cron_weekday = (dt.weekday() + 1) % 7
229 day_of_month_matches = dt.day in self.day_of_month.values
230 day_of_week_matches = cron_weekday in self.day_of_week.values
231
232 if use_or_for_days:
233 day_matches = day_of_month_matches or day_of_week_matches
234 else:
235 day_matches = day_of_month_matches and day_of_week_matches
236
237 is_valid_day = dt.month in self.month.values and day_matches
238 if is_valid_day:
239 # We're on a valid day, now find the next valid hour and minute
240 for hour in self.hour.values:
241 if hour < dt.hour:
242 continue
243 for minute in self.minute.values:
244 if hour == dt.hour and minute < dt.minute:
245 continue
246 candidate_datetime = dt.replace(hour=hour, minute=minute)
247 if candidate_datetime >= dt:
248 return candidate_datetime
249 # If no valid time is found today, reset to the first valid minute and hour of the next day
250 dt = _go_to_next_day(dt)
251 else:
252 # Increment the day until a valid month/day/weekday combination is found
253 dt = _go_to_next_day(dt)
254
255 if dt > max_future:
256 raise ValueError("No valid schedule match found in the next 500 days")
257
258
259@register_job
260class ScheduledCommand(Job):
261 """Run a shell command on a schedule."""
262
263 def __init__(self, command: str) -> None:
264 self.command = command
265
266 def __repr__(self) -> str:
267 return f"<ScheduledCommand: {self.command}>"
268
269 def run(self) -> None:
270 subprocess.run(self.command, shell=True, check=True)
271
272 def default_concurrency_key(self) -> str:
273 # The ScheduledCommand can be used for different commands,
274 # so we need the concurrency_key to separate them for uniqueness
275 return self.command
276
277
278def load_schedule(
279 schedules: list[tuple[str | Job, str | Schedule]],
280) -> list[tuple[Job, Schedule]]:
281 jobs_schedule: list[tuple[Job, Schedule]] = []
282
283 for job, schedule in schedules:
284 if isinstance(job, str):
285 if job.startswith("cmd:"):
286 job = ScheduledCommand(job[4:])
287 else:
288 job = jobs_registry.load_job(job, {"args": [], "kwargs": {}})
289
290 if isinstance(schedule, str):
291 schedule = Schedule.from_cron(schedule)
292
293 jobs_schedule.append((job, schedule))
294
295 return jobs_schedule