Draft funcy type stubs

draft-funcy-types.py

from collections.abc import Callable, Iterable, Mapping, Sequence
from re import Match
from typing import Protocol, TypeVar, overload

X = TypeVar("X")
Y = TypeVar("Y")
X_co = TypeVar("X_co", covariant=True)


class SupportsGetItem(Protocol[X_co]):
    @overload
    def __getitem__(self, key: int) -> X_co:
        ...

    @overload
    def __getitem__(self, key: slice) -> Sequence[X_co]:
        ...


SupportsGetItemX = TypeVar("SupportsGetItemX", bound=SupportsGetItem)


# ----------------------------- MAPPERS


@overload
def funcy_map(f: None, iter_: Iterable[X]) -> Iterable[X]:
    ...


@overload
def funcy_map(
    f: str, iter_: Iterable[str]
) -> Iterable[None | str | tuple[str, ...] | dict[str, str] | Match[str]]:
    ...


@overload
def funcy_map(f: int, iter_: Iterable[SupportsGetItem[X]]) -> Iterable[X]:
    ...


@overload
def funcy_map(
    f: slice, iter_: Iterable[SupportsGetItem[X]]
) -> Iterable[Sequence[X]]:
    ...


@overload
def funcy_map(
    f: Callable[[X], Y] | Mapping[X, Y], iter_: Iterable[X]
) -> Iterable[Y]:
    ...


@overload
def funcy_map(f: set[X] | frozenset[X], iter_: Iterable[Y]) -> Iterable[bool]:
    ...


# ----------------------------- FILTERS


@overload
def funcy_filter(f: None, iter_: Iterable[X]) -> Iterable[X]:
    ...


@overload
def funcy_filter(f: str, iter_: Iterable[str]) -> Iterable[str]:
    ...


@overload
def funcy_filter(
    f: int | slice, iter_: Iterable[SupportsGetItemX]
) -> Iterable[SupportsGetItemX]:
    ...


@overload
def funcy_filter(
    f: Callable[[X], bool] | Mapping[X, bool], iter_: Iterable[X]
) -> Iterable[X]:
    ...


@overload
def funcy_filter(f: set[X] | frozenset[X], iter_: Iterable[Y]) -> Iterable[Y]:
    ...