Python thenable future

ThenableFuture.py

from asyncio import Future as AsyncioFuture
from concurrent.futures import Future as ConcurrentFuture
from typing import TypeAlias

Future: TypeAlias = AsyncioFuture | ConcurrentFuture


def chain_future(first_future: Future, second_future: Future, transform_result=None):
    if first_future.done():
        exception = first_future.exception()
        if exception:
            second_future.set_exception(exception)
        else:
            result = first_future.result()
            if isinstance(result, Future):
                chain_future(result, second_future, transform_result)
            else:
                if transform_result is not None:
                    try:
                        result = transform_result(result)
                    except Exception as e:
                        second_future.set_exception(e)
                if isinstance(result, Future):
                    chain_future(result, second_future)
                else:
                    second_future.set_result(result)
    else:
        first_future.add_done_callback(
            lambda _: chain_future(first_future, second_future, transform_result)
        )


class ThenableFuture(AsyncioFuture):
    @classmethod
    def ensure_thenable(cls, future: Future):
        if isinstance(future, cls):
            return future
        new_future = cls()
        chain_future(future, new_future)
        return new_future

    @classmethod
    def resolve(cls, value):
        new_future = cls()
        if isinstance(value, Future):
            chain_future(value, new_future)
        else:
            new_future.set_result(value)
        return new_future

    def then(self, on_resolve):
        new_future = ThenableFuture()
        chain_future(self, new_future, on_resolve)
        return new_future