From 9a63a278a8665924b3c99938523dfb6ef9024b05 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 | 6 ++++ asyncio/task.py | 85 +++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 89 insertions(+), 2 deletions(-) diff --git a/asyncio/core.py b/asyncio/core.py index e31c1ce..4481c06 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.""" diff --git a/asyncio/task.py b/asyncio/task.py index 2e3a6db..ba9761c 100644 --- a/asyncio/task.py +++ b/asyncio/task.py @@ -188,7 +188,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 +211,86 @@ 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 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 + + 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('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.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)