async-chain/async_chain/chain.py

from __future__ import annotations
import inspect
import typing as t
import abc


class Chain(metaclass=abc.ABCMeta):
    __slots__ = []

    def __getattr__(self, item) -> ChainGetAttr:
        return ChainGetAttr(previous=self, item=item)

    def __getitem__(self, item) -> ChainGetItem:
        return ChainGetItem(previous=self, item=item)

    def __call__(self, *args, **kwargs) -> ChainCall:
        return ChainCall(previous=self, args=args, kwargs=kwargs)

    def __await__(self):
        return self.__evaluate__().__await__()

    @abc.abstractmethod
    def __repr__(self) -> str:
        raise NotImplementedError()

    @abc.abstractmethod
    def __display__(self) -> str:
        raise NotImplementedError()

    @abc.abstractmethod
    async def __evaluate__(self) -> t.Any:
        raise NotImplementedError()


class ChainStart(Chain):
    __slots__ = ("__start__",)

    def __init__(self, start: t.Any):
        super().__init__()
        self.__start__: t.Any = start

    def __repr__(self):
        return f"<Chain {self.__display__()}>"

    def __display__(self):
        return f"{self.__start__!r}"

    async def __evaluate__(self):
        return self.__start__


class ChainNode(Chain, metaclass=abc.ABCMeta):
    __slots__ = ("__previous__",)

    def __init__(self, previous: Chain):
        super().__init__()
        self.__previous__: Chain = previous

    def __repr__(self):
        return f"<Chain {self.__display__()}>"


class ChainGetAttr(ChainNode):
    __slots__ = ("__item__",)

    def __init__(self, previous: Chain, item: str):
        super().__init__(previous=previous)
        self.__item__: str = item

    def __display__(self):
        return f"{self.__previous__.__display__()}.{self.__item__!s}"

    async def __evaluate__(self):
        previous = await self.__previous__.__evaluate__()

        current = getattr(previous, self.__item__)

        if inspect.isawaitable(current):
            return await current
        else:
            return current


class ChainGetItem(ChainNode):
    __slots__ = ("__item__",)

    def __init__(self, previous: Chain, item: t.Any):
        super().__init__(previous=previous)
        self.__item__: t.Any = item

    def __display__(self):
        return f"{self.__previous__.__display__()}[{self.__item__!r}]"

    async def __evaluate__(self):
        previous = await self.__previous__.__evaluate__()

        current = previous[self.__item__]

        if inspect.isawaitable(current):
            return await current
        else:
            return current


class ChainCall(ChainNode):
    __slots__ = ("__args__", "__kwargs__",)

    def __init__(self, previous: Chain, args, kwargs):
        super().__init__(previous=previous)
        self.__args__ = args
        self.__kwargs__ = kwargs

    def __display__(self):
        args = map(lambda a: f"{a!r}", self.__args__)
        kwargs = map(lambda k, v: f"{k!s}={v!r}", self.__kwargs__)
        allargs = ", ".join([*args, *kwargs])
        return f"{self.__previous__.__display__()}({allargs})"

    async def __evaluate__(self):
        previous = await self.__previous__.__evaluate__()

        current = previous(*self.__args__, **self.__kwargs__)

        if inspect.isawaitable(current):
            return await current
        else:
            return current