meerschaum.utils.schedule

  1#! /usr/bin/env python3
  2"""Parse schedules and run functions at their next occurrence."""
  3
  4from __future__ import annotations
  5
  6import re
  7import threading
  8import traceback
  9from datetime import date, datetime, timezone, timedelta
 10
 11import meerschaum as mrsm
 12from meerschaum.utils.typing import Callable, Any, Optional, List, Dict
 13from meerschaum.utils.warnings import warn, error
 14
 15
 16STARTING_KEYWORD: str = 'starting'
 17INTERVAL_UNITS: List[str] = ['months', 'weeks', 'days', 'hours', 'minutes', 'seconds', 'years']
 18FREQUENCY_ALIASES: Dict[str, str] = {
 19    'daily': 'every 1 day', 'hourly': 'every 1 hour',
 20    'minutely': 'every 1 minute', 'weekly': 'every 1 week',
 21    'monthly': 'every 1 month', 'secondly': 'every 1 second',
 22    'yearly': 'every 1 year',
 23}
 24LOGIC_ALIASES: Dict[str, str] = {
 25    'and': '&', 'or': '|', ' through ': '-', ' thru ': '-', ' - ': '-',
 26}
 27CRON_DAYS_OF_WEEK: List[str] = ['mon', 'tue', 'wed', 'thu', 'fri', 'sat', 'sun']
 28CRON_DAYS_OF_WEEK_ALIASES: Dict[str, str] = {
 29    'monday': 'mon', 'tuesday': 'tue', 'tues': 'tue', 'wednesday': 'wed',
 30    'thursday': 'thu', 'thurs': 'thu', 'friday': 'fri', 'saturday': 'sat',
 31    'sunday': 'sun',
 32}
 33CRON_MONTHS: List[str] = [
 34    'jan', 'feb', 'mar', 'apr', 'may', 'jun',
 35    'jul', 'aug', 'sep', 'oct', 'nov', 'dec',
 36]
 37CRON_MONTHS_ALIASES: Dict[str, str] = {
 38    'january': 'jan', 'february': 'feb', 'march': 'mar', 'april': 'apr',
 39    'may': 'may', 'june': 'jun', 'july': 'jul', 'august': 'aug',
 40    'september': 'sep', 'october': 'oct', 'november': 'nov', 'december': 'dec',
 41}
 42SCHEDULE_ALIASES: Dict[str, str] = {
 43    **FREQUENCY_ALIASES, **LOGIC_ALIASES,
 44    **CRON_DAYS_OF_WEEK_ALIASES, **CRON_MONTHS_ALIASES,
 45}
 46
 47
 48class _IntervalTrigger:
 49    def __init__(self, start_time: datetime, **interval: float):
 50        self.start_time = start_time
 51        self._interval = timedelta(**interval)
 52        if self._interval.total_seconds() <= 0:
 53            raise ValueError("The time interval must be positive")
 54        self._last_fire_time = None
 55
 56    def next(self) -> datetime:
 57        self._last_fire_time = (
 58            self.start_time
 59            if self._last_fire_time is None
 60            else self._last_fire_time + self._interval
 61        )
 62        return self._last_fire_time
 63
 64    def next_after(self, after: datetime) -> datetime:
 65        intervals = max(0, ((after - self.start_time) // self._interval) + 1)
 66        self._last_fire_time = self.start_time + (self._interval * intervals)
 67        return self._last_fire_time
 68
 69
 70class _CalendarIntervalTrigger:
 71    """Month/year intervals which retain the original day and skip invalid dates."""
 72
 73    def __init__(self, start_time: datetime, years: int = 0, months: int = 0):
 74        if years < 0 or months < 0 or years == months == 0:
 75            raise ValueError("The calendar interval must be positive")
 76        self.start_time = start_time
 77        self.start_date = start_time.date()
 78        self.timezone = start_time.tzinfo
 79        self.years = years
 80        self.months = months
 81        self._time = start_time.timetz()
 82        self._last_fire_date = None
 83
 84    def next(self) -> datetime:
 85        previous_date = self._last_fire_date
 86        while True:
 87            if previous_date is None:
 88                next_date = self.start_date
 89            else:
 90                year, month = previous_date.year, previous_date.month
 91                while True:
 92                    month += self.months
 93                    year += self.years + ((month - 1) // 12)
 94                    month = ((month - 1) % 12) + 1
 95                    try:
 96                        next_date = date(year, month, previous_date.day)
 97                    except ValueError:
 98                        continue
 99                    break
100            candidate = datetime.fromtimestamp(
101                datetime.combine(next_date, self._time).timestamp(), self.timezone,
102            )
103            if candidate.timetz() != self._time:
104                previous_date = candidate.date()
105                continue
106            self._last_fire_date = next_date
107            return candidate
108
109    def next_after(self, after: datetime) -> datetime:
110        self._last_fire_date = None
111        candidate = self.next()
112        while candidate <= after:
113            candidate = self.next()
114        return candidate
115
116
117def _expand_cron_field(
118    expression: str,
119    minimum: int,
120    maximum: int,
121    names: Optional[List[str]] = None,
122) -> set[int]:
123    values = set()
124    names_map = {name: i + minimum for i, name in enumerate(names or [])}
125
126    def as_int(value: str) -> int:
127        return names_map.get(value.lower(), int(value) if value.lstrip('-').isdigit() else -1)
128
129    for part in expression.lower().split(','):
130        range_part, separator, step_str = part.partition('/')
131        step = int(step_str) if step_str else 1
132        if step <= 0:
133            raise ValueError(f"Invalid cron step '{step}'.")
134        if range_part == '*':
135            start, stop = minimum, maximum
136        elif '-' in range_part:
137            start_str, stop_str = range_part.split('-', 1)
138            start, stop = as_int(start_str), as_int(stop_str)
139        else:
140            start = as_int(range_part)
141            stop = maximum if separator else start
142        if start < minimum or stop > maximum or start > stop:
143            raise ValueError(f"Invalid cron field '{expression}'.")
144        values.update(range(start, stop + 1, step))
145    return values
146
147
148def _expand_cron_weekdays(expression: str) -> set[int]:
149    """Return Python weekdays while retaining cron's 0/7 = Sunday convention."""
150    values = set()
151    names_map = {
152        'sun': 0,
153        **{name: i + 1 for i, name in enumerate(CRON_DAYS_OF_WEEK[:-1])},
154    }
155
156    def as_cron_int(value: str) -> int:
157        return names_map.get(
158            value.lower(),
159            int(value) if value.isdigit() else -1,
160        )
161
162    for part in expression.lower().split(','):
163        range_part, separator, step_str = part.partition('/')
164        step = int(step_str) if separator else 1
165        if range_part == '*':
166            cron_values = range(0, 7, step)
167        else:
168            endpoints = [as_cron_int(value) for value in range_part.split('-', 1)]
169            if (
170                any(value < 0 or value > 7 for value in endpoints)
171                or step <= 0
172                or (len(endpoints) == 2 and endpoints[0] > endpoints[1])
173            ):
174                raise ValueError(f"Invalid cron field '{expression}'.")
175            cron_values = (
176                range(endpoints[0], endpoints[1] + 1, step)
177                if len(endpoints) == 2
178                else (range(endpoints[0], 8, step) if separator else endpoints)
179            )
180        if step <= 0:
181            raise ValueError(f"Invalid cron field '{expression}'.")
182        values.update(6 if value in (0, 7) else value - 1 for value in cron_values)
183    return values
184
185
186class _CronTrigger:
187    """The five-field cron subset accepted by :func:`parse_schedule`."""
188
189    def __init__(
190        self,
191        start_time: datetime,
192        minute: str = '*',
193        hour: str = '*',
194        day: str = '*',
195        month: str = '*',
196        day_of_week: str = '*',
197        year: str = '*',
198        second: str = '0',
199        microsecond: int = 0,
200    ):
201        self.start_time = start_time
202        self.timezone = start_time.tzinfo
203        self._minutes = _expand_cron_field(str(minute), 0, 59)
204        self._hours = _expand_cron_field(str(hour), 0, 23)
205        self._days = _expand_cron_field(str(day), 1, 31)
206        self._days_have_wildcard = '*' in str(day)
207        self._months = _expand_cron_field(str(month), 1, 12, CRON_MONTHS)
208        self._weekdays = _expand_cron_weekdays(str(day_of_week))
209        self._weekdays_have_wildcard = '*' in str(day_of_week)
210        self._years = None if str(year) == '*' else _expand_cron_field(str(year), 1, 9999)
211        self._seconds = _expand_cron_field(str(second), 0, 59)
212        self._microsecond = microsecond
213        self._step_seconds = 1 if len(self._seconds) > 1 else 60
214        self._last_fire_time = None
215
216    def _matches(self, candidate: datetime) -> bool:
217        return (
218            self._date_matches(candidate)
219            and candidate.hour in self._hours
220            and candidate.minute in self._minutes
221            and candidate.second in self._seconds
222        )
223
224    def _date_matches(self, candidate: datetime) -> bool:
225        day_matches = candidate.day in self._days
226        weekday_matches = candidate.weekday() in self._weekdays
227        if self._days_have_wildcard or self._weekdays_have_wildcard:
228            day_matches = day_matches and weekday_matches
229        else:
230            day_matches = day_matches or weekday_matches
231        return (
232            (self._years is None or candidate.year in self._years)
233            and candidate.month in self._months
234            and day_matches
235        )
236
237    def next(self) -> Optional[datetime]:
238        if self._last_fire_time is None:
239            candidate = self.start_time.replace(microsecond=self._microsecond)
240            if self._step_seconds == 60:
241                candidate = candidate.replace(second=min(self._seconds))
242            if candidate < self.start_time:
243                candidate = datetime.fromtimestamp(
244                    candidate.timestamp() + self._step_seconds,
245                    self.timezone,
246                )
247        else:
248            candidate = datetime.fromtimestamp(
249                self._last_fire_time.timestamp() + self._step_seconds, self.timezone,
250            )
251
252        # ponytail: skip rejected dates but keep simple stepping within an eligible day.
253        max_iterations = 11 * 366 * 24 * 60 * (60 if self._step_seconds == 1 else 1)
254        for _ in range(max_iterations):
255            if self._matches(candidate):
256                self._last_fire_time = candidate
257                return candidate
258            if not self._date_matches(candidate):
259                candidate = datetime.fromtimestamp(
260                    datetime.combine(
261                        candidate.date() + timedelta(days=1),
262                        candidate.timetz(),
263                    ).timestamp(),
264                    self.timezone,
265                )
266                continue
267            candidate = datetime.fromtimestamp(
268                candidate.timestamp() + self._step_seconds, self.timezone,
269            )
270        return None
271
272    def next_after(self, after: datetime) -> Optional[datetime]:
273        if after < self.start_time:
274            self._last_fire_time = None
275            return self.next()
276        after = after.astimezone(self.timezone)
277        candidate = after.replace(microsecond=self._microsecond)
278        if self._step_seconds == 60:
279            candidate = candidate.replace(second=min(self._seconds))
280        if candidate <= after:
281            candidate = datetime.fromtimestamp(
282                candidate.timestamp() + self._step_seconds,
283                self.timezone,
284            )
285        self._last_fire_time = datetime.fromtimestamp(
286            candidate.timestamp() - self._step_seconds,
287            self.timezone,
288        )
289        return self.next()
290
291
292class _OrTrigger:
293    def __init__(self, triggers):
294        self.triggers = triggers
295        self._next_fire_times = []
296
297    def next(self) -> Optional[datetime]:
298        if not self._next_fire_times:
299            self._next_fire_times = [trigger.next() for trigger in self.triggers]
300        earliest = min((ts for ts in self._next_fire_times if ts is not None), default=None)
301        if earliest is not None:
302            for i, fire_time in enumerate(self._next_fire_times):
303                if fire_time == earliest:
304                    self._next_fire_times[i] = self.triggers[i].next()
305        return earliest
306
307    def next_after(self, after: datetime) -> Optional[datetime]:
308        self._next_fire_times = [trigger.next_after(after) for trigger in self.triggers]
309        return self.next()
310
311
312class _AndTrigger:
313    def __init__(self, triggers, max_iterations: int = 1_000_000):
314        self.triggers = triggers
315        self.max_iterations = max_iterations
316        self._next_fire_times = []
317
318    def next(self) -> Optional[datetime]:
319        if not self._next_fire_times:
320            self._next_fire_times = [trigger.next() for trigger in self.triggers]
321        for _ in range(self.max_iterations):
322            if any(ts is None for ts in self._next_fire_times):
323                return None
324            earliest, latest = min(self._next_fire_times), max(self._next_fire_times)
325            for i, fire_time in enumerate(self._next_fire_times):
326                if fire_time == earliest:
327                    self._next_fire_times[i] = self.triggers[i].next()
328            if latest == earliest:
329                return earliest
330        raise RuntimeError("Maximum iterations reached while combining schedules.")
331
332    def next_after(self, after: datetime) -> Optional[datetime]:
333        self._next_fire_times = [trigger.next_after(after) for trigger in self.triggers]
334        return self.next()
335
336
337class _Scheduler:
338    def __init__(self):
339        self.stop_event = threading.Event()
340
341    async def stop(self):
342        self.stop_event.set()
343
344    async def wait_until_stopped(self):
345        return None
346
347
348_scheduler = None
349def schedule_function(
350    function: Callable[[Any], Any],
351    schedule: str,
352    *args,
353    debug: bool = False,
354    **kw
355) -> mrsm.SuccessTuple:
356    """Block the process and execute ``function`` according to ``schedule``."""
357    from meerschaum.utils.misc import filter_keywords
358
359    global _scheduler
360    kw['debug'] = debug
361    kw = filter_keywords(function, **kw)
362    trigger = parse_schedule(schedule, now=datetime.now(timezone.utc))
363    scheduler = _scheduler = _Scheduler()
364    pending_next_time = None
365    schedule_finished = False
366    try:
367        while not scheduler.stop_event.is_set():
368            next_time = pending_next_time or trigger.next()
369            pending_next_time = None
370            if next_time is None:
371                break
372            now = datetime.now(next_time.tzinfo or timezone.utc)
373            if next_time <= now:
374                pending_next_time = trigger.next_after(now)
375                schedule_finished = pending_next_time is None
376            if scheduler.stop_event.wait(max(0.0, (next_time - now).total_seconds())):
377                break
378            try:
379                function(*args, **kw)
380            except Exception:
381                warn(f"Scheduled function failed:\n{traceback.format_exc()}", stack=False)
382            if schedule_finished:
383                break
384    except (KeyboardInterrupt, SystemExit):
385        scheduler.stop_event.set()
386    return True, "Success"
387
388
389def parse_schedule(schedule: str, now: Optional[datetime] = None):
390    """Parse a schedule string into a stateful object with a ``next()`` method."""
391    from meerschaum.utils.misc import items_str, is_int
392
393    schedule = _canonicalize_starting_keyword(schedule)
394    starting_ts = parse_start_time(schedule, now=now)
395    schedule = schedule.split(STARTING_KEYWORD, maxsplit=1)[0].strip().lower()
396    for alias_keyword, true_keyword in SCHEDULE_ALIASES.items():
397        schedule = schedule.replace(alias_keyword, true_keyword)
398    if '&' in schedule and '|' in schedule:
399        raise ValueError("Cannot accept both 'and' + 'or' logic in the schedule frequency.")
400
401    join_str = '|' if '|' in schedule else '&'
402    schedule_parts = [part.strip() for part in schedule.split(join_str)]
403    triggers = []
404    has_seconds = 'second' in schedule
405    has_minutes = 'minute' in schedule
406    for schedule_part in schedule_parts:
407        if schedule_part.lower().startswith('every '):
408            schedule_num_str, schedule_unit = schedule_part[len('every '):].split(' ', maxsplit=1)
409            schedule_unit = schedule_unit.rstrip('s') + 's'
410            if schedule_unit not in INTERVAL_UNITS:
411                raise ValueError(
412                    f"Invalid interval '{schedule_unit}'.\n"
413                    + f"    Accepted values are {items_str(INTERVAL_UNITS)}."
414                )
415            schedule_num = int(schedule_num_str) if is_int(schedule_num_str) else float(schedule_num_str)
416            if schedule_unit in ('months', 'years'):
417                if not float(schedule_num).is_integer():
418                    raise ValueError(f"Calendar interval '{schedule_num}' must be an integer.")
419                trigger = _CalendarIntervalTrigger(
420                    starting_ts, **{schedule_unit: int(schedule_num)},
421                )
422            else:
423                trigger = _IntervalTrigger(starting_ts, **{schedule_unit: schedule_num})
424        else:
425            first_three_prefix = schedule_part[:3].lower()
426            first_four_prefix = schedule_part[:4].lower()
427            cron_kw = {}
428            if first_three_prefix in CRON_DAYS_OF_WEEK:
429                cron_kw['day_of_week'] = schedule_part
430            elif first_three_prefix in CRON_MONTHS:
431                cron_kw['month'] = schedule_part
432            elif is_int(first_four_prefix) and len(first_four_prefix) == 4:
433                cron_kw['year'] = schedule_part
434            if cron_kw:
435                trigger = _CronTrigger(
436                    starting_ts,
437                    **cron_kw,
438                    hour='*',
439                    minute='*' if (has_minutes or has_seconds) else str(starting_ts.minute),
440                    second='*' if has_seconds else str(starting_ts.second),
441                    microsecond=starting_ts.microsecond,
442                )
443            else:
444                cron_parts = schedule_part.split()
445                if len(cron_parts) != 5:
446                    raise ValueError(f"Invalid cron schedule '{schedule_part}'.")
447                trigger = _CronTrigger(
448                    starting_ts,
449                    minute=cron_parts[0], hour=cron_parts[1], day=cron_parts[2],
450                    month=cron_parts[3], day_of_week=cron_parts[4],
451                )
452        triggers.append(trigger)
453
454    if len(triggers) == 1:
455        return triggers[0]
456    return _OrTrigger(triggers) if join_str == '|' else _AndTrigger(triggers)
457
458
459def parse_start_time(schedule: str, now: Optional[datetime] = None) -> datetime:
460    """Return the explicit starting datetime in ``schedule``, or ``now``."""
461    from meerschaum.utils.dtypes import round_time
462    dateutil_parser = mrsm.attempt_import('dateutil.parser')
463    schedule = _canonicalize_starting_keyword(schedule)
464    starting_parts = schedule.split(STARTING_KEYWORD)
465    starting_str = ('now' if len(starting_parts) == 1 else starting_parts[-1]).strip()
466    now = now or datetime.now(timezone.utc)
467    try:
468        if starting_str == 'now':
469            starting_ts = now
470        elif starting_str.startswith('in '):
471            delta_vals = starting_str.replace('in ', '').split(' ', maxsplit=1)
472            delta_unit = delta_vals[-1].rstrip('s') + 's'
473            delta_num = float(delta_vals[0])
474            starting_ts = now + timedelta(**{delta_unit: delta_num})
475        elif 'tomorrow' in starting_str or 'today' in starting_str:
476            today = round_time(now, timedelta(days=1))
477            tomorrow = today + timedelta(days=1)
478            is_tomorrow = 'tomorrow' in starting_str
479            time_str = starting_str.replace('tomorrow', '').replace('today', '').strip()
480            time_ts = dateutil_parser.parse(time_str) if time_str else today
481            starting_ts = (
482                (tomorrow if is_tomorrow else today)
483                + timedelta(hours=time_ts.hour)
484                + timedelta(minutes=time_ts.minute)
485            )
486        else:
487            starting_ts = dateutil_parser.parse(starting_str)
488        schedule_parse_error = None
489    except Exception as e:
490        warn(f"Unable to parse starting time from '{starting_str}'.", stack=False)
491        schedule_parse_error = str(e)
492    if schedule_parse_error:
493        error(schedule_parse_error, ValueError, stack=False)
494    if not starting_ts.tzinfo:
495        starting_ts = starting_ts.replace(tzinfo=timezone.utc)
496    return starting_ts
497
498
499def _canonicalize_starting_keyword(schedule: str) -> str:
500    return re.sub(
501        rf'\b(?:{STARTING_KEYWORD}|beginning)\b',
502        STARTING_KEYWORD,
503        schedule,
504        flags=re.IGNORECASE,
505    )
506
507
508async def _stop_scheduler():
509    if _scheduler is None:
510        return
511    await _scheduler.stop()
512    await _scheduler.wait_until_stopped()