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 8, 2023
1 parent c277db5 commit d5f2123
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 2 deletions.
10 changes: 10 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 Expand Up @@ -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.
Expand Down
75 changes: 73 additions & 2 deletions asyncio/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
"""
Expand All @@ -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)

0 comments on commit d5f2123

Please sign in to comment.