extmod/asyncio: Add cpython task methods.

imnotjames · imnotjames · commit 15d877f96d5a · 2023-11-18T13:30:29.000-05:00
Signed-off-by: James Ward &lt;james@notjam.es&gt;
diff --git a/extmod/asyncio/core.py b/extmod/asyncio/core.py
@@ -23,6 +23,10 @@ class TimeoutError(Exception):
     pass
 
 
+class InvalidStateError(Exception):
+    pass
+
+
 # Used when calling Loop.call_exception_handler
 _exc_context = {"message": "Task exception wasn't retrieved", "exception": None, "future": None}
 
diff --git a/extmod/asyncio/task.py b/extmod/asyncio/task.py
@@ -175,3 +175,82 @@ def cancel(self):
             core._task_queue.push(self)
         self.data = core.CancelledError
         return True
+
+    def add_done_callback(self, callback):
+        if not self.state:
+            callback(self, self.data)
+
+        if self.state is not True:
+            raise RuntimeError("Tasks only support one done callback.")
+
+        self.state = callback
+
+    def remove_done_callback(self, callback):
+        if self.state is not callback:
+            return 0
+
+        self.state = True
+        return 1
+
+    def set_result(self, result):
+        raise RuntimeError()
+
+    def result(self):
+        """
+        Return the result of the Task.
+        If the Task is done, the result of the wrapped coroutine is returned (or if the coroutine raised an exception, that exception is re-raised.)
+        If the Task has been cancelled, this method raises a CancelledError exception.
+        If the Task’s result isn’t yet available, this method raises a InvalidStateError exception.
+        """
+        if not self.done():
+            raise core.InvalidStateError()
+
+        exception = self.exception()
+
+        if exception is not None:
+            raise exception
+
+        if not isinstance(self.data, StopIteration):
+            # If this isn't the case then we're in an odd state.
+            return None
+
+        return self.data.value
+
+    def set_exception(self, exception):
+        raise RuntimeError()
+
+    def exception(self):
+        """
+        Return the exception that was set on this Task.
+        The exception (or None if no exception was set) is returned only if the Task is done.
+        If the Task has been cancelled, this method raises a CancelledError exception.
+        If the Task isn’t done yet, this method raises an InvalidStateError exception.
+        """
+        if not self.done():
+            raise core.InvalidStateError()
+
+        if isinstance(self.data, core.CancelledError):
+            raise self.data
+
+        if isinstance(self.data, StopIteration):
+            # If the data is a stop iteration we can assume this
+            # was a successful run rather than any possible exception
+            return None
+
+        if not isinstance(self.data, BaseException):
+            # If the data is not any type of exception we can treat it as
+            # something else we don't understand but not an exception.
+            return None
+
+        return self.data
+
+    def cancelled(self) -> bool:
+        """
+        Return True if the Task is cancelled.
+        The Task is cancelled when the cancellation was requested with cancel() and
+        the wrapped coroutine propagated the CancelledError exception thrown into it.
+        """
+        if not self.done():
+            return False
+
+        return isinstance(self.data, core.CancelledError)
