feat: support task methods for exception and result

Author: imnotjames

asyncio/core.py

@@ -35,6 +35,12 @@ class CancelledError(BaseException):
     pass
 
 
+class InvalidStateError(Exception):
+    """Can be raised in situations like setting a result value for a task object that already has a result value set."""
+
+    pass
+
+
 class TimeoutError(Exception):
     """Raised when waiting for a task longer than the specified timeout."""
 
@@ -285,6 +291,10 @@ def run_until_complete(main_task=None):
                     _task_queue.push_head(t)
                 # Save return value of coro to pass up to caller.
                 t.data = er
+                if isinstance(er, StopIteration):
+                    t._result = er.value
+                else:
+                    t._exception = er
             elif t.state is None:
                 # Task is already finished and nothing await'ed on the task,
                 # so call the exception handler.

asyncio/task.py

@@ -145,6 +145,9 @@ class Task:
     """
 
     def __init__(self, coro, globals=None):
+        self._exception = None
+        self._result = None
+
         self.coro = coro  # Coroutine of this Task
         self.data = None  # General data for queue it is waiting on
         self.state = True  # None, False, True, a callable, or a TaskQueue instance
@@ -188,7 +191,7 @@ def done(self):
 
         return not self.state
 
-    def cancel(self):
+    def cancel(self, msg=None):
         """Cancel the task by injecting a ``CancelledError`` into it. The task
         may or may not ignore this exception.
         """
@@ -211,5 +214,73 @@ def cancel(self):
             # On the main running queue but scheduled in the future, so bring it forward to now.
             core._task_queue.remove(self)
             core._task_queue.push(self)
-        self.data = core.CancelledError
+        cancelled_error = core.CancelledError(msg)
+        self.data = cancelled_error
+        self._exception = cancelled_error
         return True
+
+    def get_coro(self):
+        return self.coro
+
+    def add_done_callback(self, callback):
+        raise NotImplementedError()
+
+    def remove_done_callback(self, callback):
+        raise NotImplementedError()
+
+    def set_result(self, result):
+        raise RuntimeError('Task does not support set_result operation')
+
+    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 InvalidStateError()
+
+        exception = self.exception()
+
+        if exception is not None:
+            raise exception
+
+        return self._result
+
+    def set_exception(self, exception):
+        raise RuntimeError('Task does not support set_exception operation')
+
+    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 InvalidStateError()
+
+        if isinstance(self._exception, core.CancelledError):
+            raise self._exception
+
+        return self._exception
+
+    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._exception, core.CancelledError)