fix tasks things

Author: DA-344

discord/client.py

@@ -235,6 +235,13 @@ def __init__(
     ):
         # self.ws is set in the connect method
         self.ws: DiscordWebSocket = None  # type: ignore
+
+        if loop is None:
+            try:
+                loop = asyncio.get_running_loop()
+            except RuntimeError:
+                pass
+
         self._loop: asyncio.AbstractEventLoop | None = loop
         self._listeners: dict[str, list[tuple[asyncio.Future, Callable[..., bool]]]] = (
             {}

discord/ext/tasks/__init__.py

@@ -28,6 +28,7 @@
 import asyncio
 import datetime
 import inspect
+import logging
 import sys
 import traceback
 from collections.abc import Sequence
@@ -43,26 +44,54 @@
 
 T = TypeVar("T")
 _func = Callable[..., Awaitable[Any]]
+_log = logging.getLogger(__name__)
 LF = TypeVar("LF", bound=_func)
 FT = TypeVar("FT", bound=_func)
 ET = TypeVar("ET", bound=Callable[[Any, BaseException], Awaitable[Any]])
 
 
+def is_ambiguous(dt: datetime.datetime) -> bool:
+    if dt.tzinfo is None or isinstance(dt.tzinfo, datetime.timezone):
+        return False
+
+    before = dt.replace(fold=0)
+    after = dt.replace(fold=1)
+
+    same_offset = before.utcoffset() == after.utcoffset()
+    same_dst = before.dst() == after.dst()
+    return not (same_offset and same_dst)
+
+
+def is_imaginary(dt: datetime.datetime) -> bool:
+    if dt.tzinfo is None or isinstance(dt.tzinfo, datetime.timezone):
+        return False
+
+    tz = dt.tzinfo
+    dt = dt.replace(tzinfo=None)
+    roundtrip = dt.replace(tzinfo=tz).astimezone(datetime.timezone.utc).astimezone(tz).replace(tzinfo=None)
+    return dt != roundtrip
+
+
 class SleepHandle:
     __slots__ = ("future", "loop", "handle")
 
     def __init__(
         self, dt: datetime.datetime, *, loop: asyncio.AbstractEventLoop
     ) -> None:
-        self.loop = loop
-        self.future = future = loop.create_future()
+        self.loop: asyncio.AbstractEventLoop = loop
+        self.future: asyncio.Future[None] = loop.create_future()
         relative_delta = discord.utils.compute_timedelta(dt)
-        self.handle = loop.call_later(relative_delta, future.set_result, True)
+        self.handle = loop.call_later(relative_delta, self._safe_result, self.future)
+
+    @staticmethod
+    def _safe_result(future: asyncio.Future) -> None:
+        if not future.done():
+            future.set_result(None)
 
     def recalculate(self, dt: datetime.datetime) -> None:
         self.handle.cancel()
         relative_delta = discord.utils.compute_timedelta(dt)
-        self.handle = self.loop.call_later(relative_delta, self.future.set_result, True)
+        self.handle = self.loop.call_later(relative_delta, self._safe_result, self.future)
 
     def wait(self) -> asyncio.Future[Any]:
         return self.future
@@ -95,7 +124,15 @@ def __init__(
     ) -> None:
         self.coro: LF = coro
         self.reconnect: bool = reconnect
-        self.loop: asyncio.AbstractEventLoop | None = loop
+
+        if loop is None:
+            try:
+                loop = asyncio.get_running_loop()
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+
+        self.loop = loop
+
         self.name: str = f'pycord-ext-task ({id(self):#x}): {coro.__qualname__}' if name in (None, MISSING) else name
         self.count: int | None = count
         self._current_loop = 0
@@ -147,53 +184,67 @@ async def _call_loop_function(self, name: str, *args: Any, **kwargs: Any) -> Non
         if name.endswith("_loop"):
             setattr(self, f"_{name}_running", False)
 
-    def _create_task(self, *args: Any, **kwargs: Any) -> asyncio.Task[None]:
-        if self.loop is None:
-            meth = asyncio.create_task
-        else:
-            meth = self.loop.create_task
-        return meth(self._loop(*args, **kwargs), name=self.name)
-
     def _try_sleep_until(self, dt: datetime.datetime):
         self._handle = SleepHandle(dt=dt, loop=asyncio.get_running_loop())
         return self._handle.wait()
 
+    def _rel_time(self) -> bool:
+        return self._time is MISSING
+
+    def _expl_time(self) -> bool:
+        return self._time is not MISSING
+
     async def _loop(self, *args: Any, **kwargs: Any) -> None:
         backoff = ExponentialBackoff()
         await self._call_loop_function("before_loop")
         self._last_iteration_failed = False
-        if self._time is not MISSING:
-            # the time index should be prepared every time the internal loop is started
-            self._prepare_time_index()
+        if self._expl_time():
             self._next_iteration = self._get_next_sleep_time()
         else:
             self._next_iteration = datetime.datetime.now(datetime.timezone.utc)
+
         try:
-            await self._try_sleep_until(self._next_iteration)
+            if self._stop_next_iteration:
+                return
+
             while True:
+                if self._expl_time():
+                    await self._try_sleep_until(self._next_iteration)
                 if not self._last_iteration_failed:
                     self._last_iteration = self._next_iteration
                     self._next_iteration = self._get_next_sleep_time()
+
+                    while self._expl_time() and self._next_iteration <= self._last_iteration:
+                        _log.warning(
+                            'Task %s woke up at %s, which was before expected (%s). Sleeping again to fix it...',
+                            self.coro.__name__,
+                            discord.utils.utcnow(),
+                            self._next_iteration,
+                        )
+                        await self._try_sleep_until(self._next_iteration)
+                        self._next_iteration = self._get_next_sleep_time()
                 try:
                     await self.coro(*args, **kwargs)
                     self._last_iteration_failed = False
-                    backoff = ExponentialBackoff()
-                except self._valid_exception:
+                except self._valid_exception as exc:
                     self._last_iteration_failed = True
                     if not self.reconnect:
                         raise
-                    await asyncio.sleep(backoff.delay())
-                else:
-                    await self._try_sleep_until(self._next_iteration)
 
+                    delay = backoff.delay()
+                    _log.warning(
+                        'Received an exception which was in the valid exception set. Task will run again in %s.2f seconds',
+                        self.coro.__name__,
+                        delay,
+                        exc_info=exc,
+                    )
+                    await asyncio.sleep(delay)
+                else:
                     if self._stop_next_iteration:
                         return
 
-                    now = datetime.datetime.now(datetime.timezone.utc)
-                    if now > self._next_iteration:
-                        self._next_iteration = now
-                        if self._time is not MISSING:
-                            self._prepare_time_index(now)
+                    if self._rel_time():
+                        await self._try_sleep_until(self._next_iteration)
 
                     self._current_loop += 1
                     if self._current_loop == self.count:
@@ -208,7 +259,8 @@ async def _loop(self, *args: Any, **kwargs: Any) -> None:
             raise exc
         finally:
             await self._call_loop_function("after_loop")
-            self._handle.cancel()
+            if self._handle:
+                self._handle.cancel()
             self._is_being_cancelled = False
             self._current_loop = 0
             self._stop_next_iteration = False
@@ -226,8 +278,8 @@ def __get__(self, obj: T, objtype: type[T]) -> Loop[LF]:
             time=self._time,
             count=self.count,
             reconnect=self.reconnect,
-            loop=self.loop,
             name=self.name,
+            loop=self.loop,
         )
         copy._injected = obj
         copy._before_loop = self._before_loop
@@ -340,7 +392,7 @@ def start(self, *args: Any, **kwargs: Any) -> asyncio.Task[None]:
         if self._injected is not None:
             args = (self._injected, *args)
 
-        self._task = self._create_task(*args, **kwargs)
+        self._task = self.loop.create_task(self._loop(*args, **kwargs), name=self.name)
         return self._task
 
     def stop(self) -> None:
@@ -574,66 +626,51 @@ def error(self, coro: ET) -> ET:
         self._error = coro  # type: ignore
         return coro
 
-    def _get_next_sleep_time(self) -> datetime.datetime:
+    def _get_next_sleep_time(self, now: datetime.datetime = MISSING) -> datetime.datetime:
         if self._sleep is not MISSING:
             return self._last_iteration + datetime.timedelta(seconds=self._sleep)
 
-        if self._time_index >= len(self._time):
-            self._time_index = 0
-            if self._current_loop == 0:
-                # if we're at the last index on the first iteration, we need to sleep until tomorrow
-                return datetime.datetime.combine(
-                    datetime.datetime.now(self._time[0].tzinfo or datetime.timezone.utc)
-                    + datetime.timedelta(days=1),
-                    self._time[0],
-                )
+        if now is MISSING:
+            now = datetime.datetime.now(datetime.timezone.utc)
 
-        next_time = self._time[self._time_index]
-
-        if self._current_loop == 0:
-            self._time_index += 1
-            if (
-                next_time
-                > datetime.datetime.now(
-                    next_time.tzinfo or datetime.timezone.utc
-                ).timetz()
-            ):
-                return datetime.datetime.combine(
-                    datetime.datetime.now(next_time.tzinfo or datetime.timezone.utc),
-                    next_time,
-                )
-            else:
-                return datetime.datetime.combine(
-                    datetime.datetime.now(next_time.tzinfo or datetime.timezone.utc)
-                    + datetime.timedelta(days=1),
-                    next_time,
-                )
+        index = self._start_time_relative_to(now)
 
-        next_date = cast(
-            datetime.datetime, self._last_iteration.astimezone(next_time.tzinfo)
-        )
-        if next_time < next_date.timetz():
-            next_date += datetime.timedelta(days=1)
+        if index is None:
+            time = self._time[0]
+            tomorrow = now.astimezone(time.tzinfo) + datetime.timedelta(days=1)
+            date = tomorrow.date()
+        else:
+            time = self._time[index]
+            date = now.astimezone(time.tzinfo).date()
+
+        dt = datetime.datetime.combine(date, time, tzinfo=time.tzinfo)
 
-        self._time_index += 1
-        return datetime.datetime.combine(next_date, next_time)
+        if dt.tzinfo is None or isinstance(dt.tzinfo, datetime.timezone):
+            return dt
+
+        if is_imaginary(dt):
+            tomorrow = dt + datetime.timedelta(days=1)
+            yesterday = dt - datetime.timedelta(days=1)
+            return dt + (tomorrow.utcoffset() - yesterday.utcoffset())  # type: ignore
+        elif is_ambiguous(dt):
+            return dt.replace(fold=1)
+        else:
+            return dt
 
-    def _prepare_time_index(self, now: datetime.datetime = MISSING) -> None:
+    def _start_time_relative_to(self, now: datetime.datetime) -> int | None:
         # now kwarg should be a datetime.datetime representing the time "now"
         # to calculate the next time index from
 
         # pre-condition: self._time is set
-        time_now = (
-            now
-            if now is not MISSING
-            else datetime.datetime.now(datetime.timezone.utc).replace(microsecond=0)
-        )
         for idx, time in enumerate(self._time):
-            if time >= time_now.astimezone(time.tzinfo).timetz():
-                self._time_index = idx
-                break
+            # Convert the current time to the target timezone
+            # e.g. 18:00 UTC -> 03:00 UTC+9
+            # Then compare the time instances to see if they're the same
+            start = now.astimezone(time.tzinfo)
+            if time >= start.timetz():
+                return idx
         else:
-            self._time_index = 0
+            return None
 
     def _get_time_parameter(
         self,
@@ -780,9 +817,6 @@ def loop(
         one used in :meth:`discord.Client.connect`.
     loop: Optional[:class:`asyncio.AbstractEventLoop`]
         The loop to use to register the task, defaults to ``None``.
-
-        .. versionchanged:: 2.7
-            This can now be ``None``
     name: Optional[:class:`str`]
         The name to create the task with, defaults to ``None``.
 
@@ -806,8 +840,8 @@ def decorator(func: LF) -> Loop[LF]:
             count=count,
             time=time,
             reconnect=reconnect,
-            loop=loop,
             name=name,
+            loop=loop,
         )
 
     return decorator