From 0c458f4469fe834750ee44146e869e0fe5588d9a Mon Sep 17 00:00:00 2001
From: Hood Chatham <roberthoodchatham@gmail.com>
Date: Mon, 10 Apr 2023 22:09:06 -0700
Subject: [PATCH] ENH Add then, catch, and finally_ to the Tasks too (#3748)

---
 docs/project/changelog.md |  4 ++++
 src/py/pyodide/webloop.py | 27 ++++++++++++++++-------
 src/tests/test_webloop.py | 45 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 8 deletions(-)

diff --git a/docs/project/changelog.md b/docs/project/changelog.md
index e85cfec6c..10d80044d 100644
--- a/docs/project/changelog.md
+++ b/docs/project/changelog.md
@@ -15,6 +15,10 @@ myst:
 
 ## Unreleased
 
+- {{ Enhancement }} The promise methods `then`, `catch` and `finally_` are now
+  present also on `Task`s as well as `Future`s.
+  {pr}`3748`
+
 ### Deployment
 
 - {{ Fix }} Export `python_stdlib.zip` in `package.json`.
diff --git a/src/py/pyodide/webloop.py b/src/py/pyodide/webloop.py
index 5fda5f5b2..7ebbeabfe 100644
--- a/src/py/pyodide/webloop.py
+++ b/src/py/pyodide/webloop.py
@@ -4,7 +4,7 @@ import inspect
 import sys
 import time
 import traceback
-from asyncio import Future
+from asyncio import Future, Task
 from collections.abc import Awaitable, Callable
 from typing import Any, TypeVar, overload
 
@@ -18,10 +18,11 @@ S = TypeVar("S")
 
 
 class PyodideFuture(Future[T]):
-    """A future with extra ``then``, ``catch``, and ``finally_`` methods based
-    on the Javascript promise API. ``Webloop.create_future`` returns these so in
-    practice all futures encountered in Pyodide should be an instance of
-    ``PyodideFuture``.
+    """A :py:class:`~asyncio.Future` with extra :js:meth:`~Promise.then`,
+    :js:meth:`~Promise.catch`, and :js:meth:`finally_() <Promise.finally>` methods
+    based on the Javascript promise API. :py:meth:`~asyncio.loop.create_future`
+    returns these so in practice all futures encountered in Pyodide should be an
+    instance of :py:class:`~pyodide.webloop.PyodideFuture`.
     """
 
     @overload
@@ -141,7 +142,7 @@ class PyodideFuture(Future[T]):
         return self.then(None, onrejected)
 
     def finally_(self, onfinally: Callable[[], None]) -> "PyodideFuture[T]":
-        """When the future is either resolved or rejected, call onfinally with
+        """When the future is either resolved or rejected, call ``onfinally`` with
         no arguments.
         """
         result: PyodideFuture[T] = PyodideFuture()
@@ -167,6 +168,16 @@ class PyodideFuture(Future[T]):
         return result
 
 
+class PyodideTask(Task[T], PyodideFuture[T]):
+    """Inherits from both :py:class:`~asyncio.Task` and
+    :py:class:`~pyodide.webloop.PyodideFuture`
+
+    Instantiation is discouraged unless you are writing your own event loop.
+    """
+
+    pass
+
+
 class WebLoop(asyncio.AbstractEventLoop):
     """A custom event loop for use in Pyodide.
 
@@ -400,7 +411,7 @@ class WebLoop(asyncio.AbstractEventLoop):
         """
         self._check_closed()
         if self._task_factory is None:
-            task = asyncio.tasks.Task(coro, loop=self, name=name)
+            task = PyodideTask(coro, loop=self, name=name)
             if task._source_traceback:  # type: ignore[attr-defined]
                 # Added comment:
                 # this only happens if get_debug() returns True.
@@ -599,4 +610,4 @@ def _initialize_event_loop():
     policy.get_event_loop()
 
 
-__all__ = ["WebLoop", "WebLoopPolicy", "PyodideFuture"]
+__all__ = ["WebLoop", "WebLoopPolicy", "PyodideFuture", "PyodideTask"]
diff --git a/src/tests/test_webloop.py b/src/tests/test_webloop.py
index 74ef887a6..fa3e663fd 100644
--- a/src/tests/test_webloop.py
+++ b/src/tests/test_webloop.py
@@ -355,6 +355,51 @@ async def test_pyodide_future2(selenium):
     assert name == "pytest"
 
 
+@run_in_pyodide
+async def test_pyodide_task(selenium):
+    from asyncio import Future, ensure_future, sleep
+
+    async def taskify(fut):
+        return await fut
+
+    def do_the_thing():
+        d = dict(
+            did_onresolve=None,
+            did_onreject=None,
+            did_onfinally=False,
+        )
+        f: Future[int] = Future()
+        t = ensure_future(taskify(f))
+        t.then(
+            lambda v: d.update(did_onresolve=v), lambda e: d.update(did_onreject=e)
+        ).finally_(lambda: d.update(did_onfinally=True))
+        return f, d
+
+    f, d = do_the_thing()
+    f.set_result(7)
+    await sleep(0.1)
+    assert d == dict(
+        did_onresolve=7,
+        did_onreject=None,
+        did_onfinally=True,
+    )
+
+    f, d = do_the_thing()
+    e = Exception("Oops!")
+    f.set_exception(e)
+    assert d == dict(
+        did_onresolve=None,
+        did_onreject=None,
+        did_onfinally=False,
+    )
+    await sleep(0.1)
+    assert d == dict(
+        did_onresolve=None,
+        did_onreject=e,
+        did_onfinally=True,
+    )
+
+
 @run_in_pyodide
 async def test_inprogress(selenium):
     import asyncio