Support for async/await in our own Implementation

So what’s necessary for supporting async/await in our implementation? Just a single line needs to be added:

Source

Future v5

from enum import Enum
from typing import Any, Callable, Generator

from loop import Loop


class CancelledError(Exception):
    """
    Raised if a Future is cancelled
    """


class FutureState(Enum):
    PENDING = "pending"
    DONE = "done"
    CANCELLED = "cancelled"


class Future:
    """Return a result in the future v5"""

    _result = None

    def __init__(self, name: str = None):
        self._name = name
        self._callbacks = []
        self._state = FutureState.PENDING
        self._loop = Loop.get_current_loop()

    def set_result(self, result: Any):
        if self._state != FutureState.PENDING:
            raise RuntimeError("Invalid Future state")

        self._result = result
        self._state = FutureState.DONE
        self._schedule_callbacks()

    def cancel(self) -> bool:
        if self._state != FutureState.PENDING:
            return False

        self._state = FutureState.CANCELLED
        self._schedule_callbacks()
        return True

    def result(self) -> Any:
        if self._state == FutureState.CANCELLED:
            raise CancelledError()
        if self._state != FutureState.DONE:
            raise RuntimeError("Invalid Future state")
        return self._result

    def done(self) -> bool:
        return self._state == FutureState.DONE

    def cancelled(self) -> bool:
        return self._state == FutureState.CANCELLED

    def add_done_callback(self, fn: Callable[["Future"], None]) -> None:
        if self._state != FutureState.PENDING:
            # we already have a result or are cancelled
            self._loop.schedule(self._name, fn, self)
        else:
            self._callbacks.append(fn)

    def _schedule_callbacks(self) -> None:
        if not self._callbacks:
            return

        callbacks = self._callbacks.copy()
        self._callbacks.clear()

        for callback in callbacks:
            self._loop.schedule(self._name, callback, self)

    def __repr__(self) -> str:
        return (
            f"<{self.__class__.__name__} name='{self._name}' "
            f"state={self._state} id='{hex(id(self))}'>"
        )

    def __iter__(self) -> Generator["Future", None, Any]:
        yield self
        return self.result()

    __await__ = __iter__

Diff

Future v5

--- future4.py	2022-11-09 12:22:03.254019054 +0100
+++ future5.py	2022-11-16 12:18:52.499094129 +0100
@@ -17,7 +17,7 @@
 
 
 class Future:
-    """Return a result in the future v4"""
+    """Return a result in the future v5"""
 
     _result = None
 
@@ -82,3 +82,5 @@
     def __iter__(self) -> Generator["Future", None, Any]:
         yield self  # some new magic
         return self.result()
+
+    __await__ = __iter__

from future import Future
from loop import Loop
from task import Task


async def some_result(result):
    future = Future("Some Result")
    future.set_result(result)
    return await future


async def add(coroutine1, coroutine2):
    task1 = Task(coroutine1, "Add X")
    task2 = Task(coroutine2, "Add Y")
    x = await task1
    y = await task2
    return x + y


async def main():
    return await add(some_result(1), some_result(2))


loop = Loop.get_current_loop()
result = loop.run(main())
print("Loop finished with result", result)

Output:

Loop step 1 [<Handle name='Initial Task' callback='step'>]
Loop step 2 [<Handle name='Add X' callback='step'>, <Handle name='Add Y' callback='step'>]
Loop step 3 [<Handle name='Some Result' callback='_wakeup'>, <Handle name='Some Result' callback='_wakeup'>]
Loop step 4 [<Handle name='Add X' callback='step'>, <Handle name='Add Y' callback='step'>]
Loop step 5 [<Handle name='Add X' callback='_wakeup'>]
Loop step 6 [<Handle name='Initial Task' callback='step'>]
Loop step 7 [<Handle name='Add Y' callback='_wakeup'>]
Loop step 8 [<Handle name='Initial Task' callback='step'>]
Loop step 9 [<Handle name='Initial Task' callback='_done'>]
Loop finished with result 3