diff --git a/docs/reference/mode.utils.cron.rst b/docs/reference/mode.utils.cron.rst new file mode 100644 index 00000000..4f02440c --- /dev/null +++ b/docs/reference/mode.utils.cron.rst @@ -0,0 +1,11 @@ +===================================================== + ``mode.utils.cron`` +===================================================== + +.. contents:: + :local: +.. currentmodule:: mode.utils.cron + +.. automodule:: mode.utils.cron + :members: + :undoc-members: diff --git a/mode/services.py b/mode/services.py index 5a7b31b1..fb5e774a 100644 --- a/mode/services.py +++ b/mode/services.py @@ -5,6 +5,7 @@ from functools import wraps from time import monotonic, perf_counter +from datetime import tzinfo from types import TracebackType from typing import ( Any, @@ -36,6 +37,7 @@ from .utils.objects import iter_mro_reversed, qualname from .utils.text import maybecat from .utils.times import Seconds, want_seconds +from .utils.cron import secs_for_next from .utils.tracebacks import format_task_stack from .utils.trees import Node from .utils.typing import AsyncContextManager @@ -447,6 +449,34 @@ async def _repeater(self: Service) -> None: return cls.task(_repeater) return _decorate + @classmethod + def crontab(cls, cron_format: str, *, + timezone: tzinfo = None) -> Callable[[Callable], ServiceTask]: + """Background timer executing periodic task based on Crontab description. + + Example: + >>> class S(Service): + ... + ... @Service.crontab(cron_format='30 18 * * *', + timezone=pytz.timezone('US/Pacific')) + ... async def every_6_30_pm_pacific(self): + ... print('IT IS 6:30pm') + ... + ... @Service.crontab(cron_format='30 18 * * *') + ... async def every_6_30_pm(self): + ... print('6:30pm UTC') + """ + def _decorate( + fun: Callable[[ServiceT], Awaitable[None]]) -> ServiceTask: + @wraps(fun) + async def _cron_starter(self: Service) -> None: + while not self.should_stop: + await self.sleep(secs_for_next(cron_format, timezone)) + if not self.should_stop: + await fun(self) + return cls.task(_cron_starter) + return _decorate + @classmethod def transitions_to(cls, flag: str) -> Callable: """Decorate function to set and reset diagnostic flag.""" diff --git a/mode/utils/cron.py b/mode/utils/cron.py new file mode 100644 index 00000000..b2fcf638 --- /dev/null +++ b/mode/utils/cron.py @@ -0,0 +1,17 @@ +"""Crontab Utilities.""" +import time +from typing import cast +from datetime import datetime, tzinfo +from croniter.croniter import croniter + + +def secs_for_next(cron_format: str, tz: tzinfo = None) -> float: + """Return seconds until next execution given Crontab style format.""" + now_ts = time.time() + # If we have a tz object we'll make now timezone aware, and + # if not will set now to be the current timestamp (tz + # unaware) + # If we have tz, now will be a datetime, if not an integer + now = tz and datetime.now(tz) or now_ts + cron_it = croniter(cron_format, start_time=now) + return cast(float, cron_it.get_next(float)) - now_ts diff --git a/requirements/default.txt b/requirements/default.txt index 8e4ac51f..a3ce5231 100644 --- a/requirements/default.txt +++ b/requirements/default.txt @@ -2,3 +2,4 @@ colorlog>=2.9.0 aiocontextvars>=0.2 ; python_version<'3.7' mypy_extensions typing_extensions; python_version<'3.8' +croniter>=0.3.16 diff --git a/requirements/test.txt b/requirements/test.txt index 35b0bd75..72213c4c 100644 --- a/requirements/test.txt +++ b/requirements/test.txt @@ -1,4 +1,5 @@ hypothesis>=3.31 +freezegun>=0.3.11 pytest-aiofiles>=0.2.0 pytest-asyncio>=0.8 pytest-base-url>=1.4.1 diff --git a/t/unit/test_services.py b/t/unit/test_services.py index 731088a2..05110e79 100644 --- a/t/unit/test_services.py +++ b/t/unit/test_services.py @@ -1,4 +1,5 @@ import asyncio +from time import sleep from typing import ContextManager from mode import Service from mode.services import Diag, ServiceTask, WaitResult @@ -218,6 +219,27 @@ async def itertimer(*args, **kwargs): pass assert m.call_count == 4 + @pytest.mark.asyncio + async def test_crontab(self): + m = Mock() + + with patch('mode.services.secs_for_next') as secs_for_next: + secs_for_next.secs_for_next.return_value = 0.1 + + class Foo(Service): + + @Service.crontab('* * * * *') + async def foo(self): + m() + self._stopped.set() + + foo = Foo() + foo.sleep = AsyncMock() + async with foo: + await asyncio.sleep(0) + + m.assert_called_once_with() + @pytest.mark.asyncio async def test_transitions_to(self, *, service): diff --git a/t/unit/utils/test_cron.py b/t/unit/utils/test_cron.py new file mode 100644 index 00000000..49aa3e19 --- /dev/null +++ b/t/unit/utils/test_cron.py @@ -0,0 +1,30 @@ +import pytz +from freezegun import freeze_time +from mode.utils.cron import secs_for_next + +SECS_IN_HOUR = 60 * 60 + + +@freeze_time('2000-01-01 00:00:00') +def test_secs_for_next(): + every_minute_cron_format = '*/1 * * * *' + assert secs_for_next(every_minute_cron_format) == 60 + + every_8pm_cron_format = '0 20 * * *' + assert secs_for_next(every_8pm_cron_format) == 20 * SECS_IN_HOUR + + every_4th_july_1pm_cron_format = '0 13 4 7 *' + days_until_4th_july = 31 + 28 + 31 + 30 + 31 + 30 + 4 + secs_until_4th_july = SECS_IN_HOUR * 24 * days_until_4th_july + secs_until_1_pm = 13 * SECS_IN_HOUR + total_secs = secs_until_4th_july + secs_until_1_pm + assert secs_for_next(every_4th_july_1pm_cron_format) == total_secs + + +@freeze_time('2000-01-01 00:00:00') +def test_secs_for_next_with_tz(): + pacific = pytz.timezone('US/Pacific') + + every_8pm_cron_format = '0 20 * * *' + # In Pacific time it's 16:00 so only 4 hours until 8:00pm + assert secs_for_next(every_8pm_cron_format, tz=pacific) == 4 * SECS_IN_HOUR