polymorphic python fmap

fmap.py

from __future__ import annotations

from asyncio import Task, create_task
from dataclasses import dataclass
from typing import (
    Any,
    AsyncIterable,
    AsyncIterator,
    Awaitable,
    Callable,
    Generic,
    Iterable,
    Iterator,
    ParamSpec,
    TypeVar,
    overload,
)

_A = TypeVar("_A")
_B = TypeVar("_B")
_K = TypeVar("_K")
_P = ParamSpec("_P")


@dataclass(frozen=True)
class fmap(Generic[_A, _B]):
    f: Callable[[_A], _B]

    # fmt: off
    # # my adts
    # @overload
    # def __call__(self, fa: Result[_A, _E]) -> Result[_B, _E]: ...
    # @overload
    # def __call__(self, fa: Lsd[_A]) -> Lsd[_B]: ...
    # @overload
    # def __call__(self, fa: Tree[_A]) ->  Tree[_B]: ...

    # async
    @overload
    def __call__(self, fa: Task[_A]) ->  Task[_B]: ...
    @overload
    def __call__(self, fa: Awaitable[_A]) ->  Awaitable[_B]: ...
    @overload
    def __call__(self, fa: AsyncIterator[_A]) ->  AsyncIterator[_B]: ...
    @overload
    def __call__(self, fa: AsyncIterable[_A]) ->  AsyncIterable[_B]: ...

    # containers
    @overload
    def __call__(self, fa: list[_A]) ->  list[_B]: ...
    @overload
    def __call__(self, fa: dict[_K, _A]) ->  dict[_K, _B]: ...
    @overload
    def __call__(self, fa: Iterator[_A]) ->  Iterator[_B]: ...
    @overload
    def __call__(self, fa: Iterable[_A]) ->  Iterable[_B]: ...
    @overload
    def __call__(self, fa: Callable[_P, _A]) ->  Callable[_P, _B]: ...
    # fmt: on

    def __call__(self, fa: Any) -> Any:
        f = self.f
        try:
            fa.fmap  # Do I have my own fmap method?
            # if so, don't use here, do it down there as to be sure
            # the 'AttributeError' happened here and not in `f`
        except AttributeError:
            # fmt: off
            if isinstance(fa, Task):
                async def _(): return f(await fa) # type: ignore
                return create_task(_())
            if isinstance(fa, Awaitable):
                async def _(): return f(await fa) # type: ignore
                return create_task(_())
            if isinstance(fa, AsyncIterator): return (f(a) async for a in fa) # type: ignore
            if isinstance(fa, AsyncIterable):
                class aiterable: # we need to wrap it in a class to make it reusable
                    def __aiter__(self): return (f(a) async for a in fa)
                return aiterable()


            if isinstance(fa, list): return [f(a) for a in fa] # type: ignore
            if isinstance(fa, dict): return {k: f(v) for k, v in fa.items()} # type: ignore
            if isinstance(fa, Iterator): return (f(a) for a in fa) # type: ignore
            if isinstance(fa, Iterable): return (f(a) for a in fa) # type: ignore
            if isinstance(fa, Callable): return lambda *a, **kwa: f(fa(*a, **kwa)) # type: ignore
            # fmt: on
        return fa.map(f)