From d5f21238347367dc302f5975511edcc0c6573b92 Mon Sep 17 00:00:00 2001 From: James Ward Date: Wed, 8 Nov 2023 16:42:34 -0500 Subject: [PATCH] feat: support task methods for `exception` and `result` --- asyncio/core.py | 10 +++++++ asyncio/task.py | 75 +++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/asyncio/core.py b/asyncio/core.py index e31c1ce..6769c12 100644 --- a/asyncio/core.py +++ b/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. diff --git a/asyncio/task.py b/asyncio/task.py index 2e3a6db..8e705ae 100644 --- a/asyncio/task.py +++ b/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)