Skip to content

Commit

Permalink
feat: support task methods for exception and result
Browse files Browse the repository at this point in the history
  • Loading branch information
imnotjames committed Nov 9, 2023
1 parent c277db5 commit 9a63a27
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 2 deletions.
6 changes: 6 additions & 0 deletions asyncio/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down
85 changes: 83 additions & 2 deletions asyncio/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand All @@ -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)

0 comments on commit 9a63a27

Please sign in to comment.