diff --git a/discord/ext/tasks/__init__.py b/discord/ext/tasks/__init__.py index 6a7c9437c..dbbd2df78 100644 --- a/discord/ext/tasks/__init__.py +++ b/discord/ext/tasks/__init__.py @@ -4,19 +4,36 @@ import websockets import discord import inspect import logging +import datetime +from collections.abc import Sequence from discord.backoff import ExponentialBackoff MAX_ASYNCIO_SECONDS = 3456000 log = logging.getLogger(__name__) +def _get_time_parameter(time, *, inst=isinstance, dt=datetime.time, utc=datetime.timezone.utc): + if inst(time, dt): + return [time if time.tzinfo is not None else time.replace(tzinfo=utc)] + if not inst(time, Sequence): + raise TypeError('time parameter must be datetime.time or a sequence of datetime.time') + if not time: + raise ValueError('time parameter must not be an empty sequence.') + + ret = [] + for index, t in enumerate(time): + if not inst(t, dt): + raise TypeError('index %d of time sequence expected %r not %r' % (index, dt, type(t))) + ret.append(t if t.tzinfo is not None else t.replace(tzinfo=utc)) + return sorted(ret) + class Loop: """A background task helper that abstracts the loop and reconnection logic for you. The main interface to create this is through :func:`loop`. """ - def __init__(self, coro, seconds, hours, minutes, count, reconnect, loop): + def __init__(self, coro, seconds, hours, minutes, time, count, reconnect, loop): self.coro = coro self.reconnect = reconnect self.loop = loop or asyncio.get_event_loop() @@ -45,7 +62,7 @@ class Loop: if self.count is not None and self.count <= 0: raise ValueError('count must be greater than 0 or None.') - self.change_interval(seconds=seconds, minutes=minutes, hours=hours) + self.change_interval(seconds=seconds, minutes=minutes, hours=hours, time=time) if not inspect.iscoroutinefunction(self.coro): raise TypeError('Expected coroutine function, not {0.__name__!r}.'.format(type(self.coro))) @@ -64,6 +81,10 @@ class Loop: backoff = ExponentialBackoff() await self._call_loop_function('before_loop') try: + # If a specific time is needed, wait before calling the function + if self._time is not None: + await asyncio.sleep(self._get_next_sleep_time()) + while True: try: await self.coro(*args, **kwargs) @@ -78,7 +99,7 @@ class Loop: if self._current_loop == self.count: break - await asyncio.sleep(self._sleep) + await asyncio.sleep(self._get_next_sleep_time()) except asyncio.CancelledError: self._is_being_cancelled = True raise @@ -321,7 +342,38 @@ class Loop: self._after_loop = coro return coro - def change_interval(self, *, seconds=0, minutes=0, hours=0): + def _get_next_sleep_time(self): + if self._sleep is not None: + return self._sleep + + # microseconds in the calculations sometimes leads to the sleep time + # being too small + now = datetime.datetime.now(datetime.timezone.utc).replace(microsecond=0) + if self._time_index >= len(self._time): + self._time_index = 0 + + # note: self._time is sorted by earliest -> latest + current_time = self._time[self._time_index] + if current_time >= now.timetz(): + as_dt = datetime.datetime.combine(now.date(), current_time) + else: + tomorrow = now + datetime.timedelta(days=1) + as_dt = datetime.datetime.combine(tomorrow.date(), current_time) + + delta = (as_dt - now).total_seconds() + self._time_index += 1 + return max(delta, 0.0) + + def _prepare_index(self): + # pre-condition: self._time is set + # find the current index that we should be in + now = datetime.datetime.now(datetime.timezone.utc).replace(microsecond=0).timetz() + for index, dt in enumerate(self._time): + if dt >= now: + self._time_index = index + break + + def change_interval(self, *, seconds=0, minutes=0, hours=0, time=None): """Changes the interval for the sleep time. .. note:: @@ -339,27 +391,47 @@ class Loop: The number of minutes between every iteration. hours: :class:`float` The number of hours between every iteration. + time: Union[Sequence[:class:`datetime.time`], :class:`datetime.time`] + The exact times to run this loop at. Either a list or a single + value of :class:`datetime.time` should be passed. Note that + this cannot be mixed with the relative time parameters. + + .. versionadded:: 1.3.0 Raises ------- ValueError An invalid value was given. + TypeError + Mixing ``time`` parameter with relative time parameter or + passing an improper type for the ``time`` parameter. """ - sleep = seconds + (minutes * 60.0) + (hours * 3600.0) - if sleep >= MAX_ASYNCIO_SECONDS: - fmt = 'Total number of seconds exceeds asyncio imposed limit of {0} seconds.' - raise ValueError(fmt.format(MAX_ASYNCIO_SECONDS)) + if any((seconds, minutes, hours)) and time is not None: + raise TypeError('Cannot mix relative time with explicit time.') - if sleep < 0: - raise ValueError('Total number of seconds cannot be less than zero.') - - self._sleep = sleep self.seconds = seconds self.hours = hours self.minutes = minutes + self._time_index = 0 + + if time is None: + sleep = seconds + (minutes * 60.0) + (hours * 3600.0) + if sleep >= MAX_ASYNCIO_SECONDS: + fmt = 'Total number of seconds exceeds asyncio imposed limit of {0} seconds.' + raise ValueError(fmt.format(MAX_ASYNCIO_SECONDS)) -def loop(*, seconds=0, minutes=0, hours=0, count=None, reconnect=True, loop=None): + if sleep < 0: + raise ValueError('Total number of seconds cannot be less than zero.') + + self._sleep = sleep + self._time = None + else: + self._sleep = None + self._time = _get_time_parameter(time) + self._prepare_index() + +def loop(*, seconds=0, minutes=0, hours=0, count=None, time=None, reconnect=True, loop=None): """A decorator that schedules a task in the background for you with optional reconnect logic. The decorator returns a :class:`Loop`. @@ -374,6 +446,12 @@ def loop(*, seconds=0, minutes=0, hours=0, count=None, reconnect=True, loop=None count: Optional[:class:`int`] The number of loops to do, ``None`` if it should be an infinite loop. + time: Union[Sequence[:class:`datetime.time`], :class:`datetime.time`] + The exact times to run this loop at. Either a list or a single + value of :class:`datetime.time` should be passed. Note that + this cannot be mixed with the relative time parameters. + + .. versionadded:: 1.3.0 reconnect: :class:`bool` Whether to handle errors and restart the task using an exponential back-off algorithm similar to the @@ -391,5 +469,5 @@ def loop(*, seconds=0, minutes=0, hours=0, count=None, reconnect=True, loop=None """ def decorator(func): return Loop(func, seconds=seconds, minutes=minutes, hours=hours, - count=count, reconnect=reconnect, loop=loop) + time=time, count=count, reconnect=reconnect, loop=loop) return decorator