diff --git a/pipe.py b/pipe.py index 1af26e4..f8884f1 100644 --- a/pipe.py +++ b/pipe.py @@ -12,6 +12,7 @@ import sys from collections import deque from contextlib import closing +import typing class Pipe: @@ -251,6 +252,12 @@ def batched(iterable, n): yield batch +@Pipe +def cast(iterable, type_): + """Cast an iterable to a given type.""" + return typing.cast(type_, iterable) + + chain = Pipe(itertools.chain.from_iterable) chain_with = Pipe(itertools.chain) islice = Pipe(itertools.islice) diff --git a/pipe.pyi b/pipe.pyi new file mode 100644 index 0000000..b65f1d0 --- /dev/null +++ b/pipe.pyi @@ -0,0 +1,410 @@ +from collections.abc import Sequence +from typing import ( + Optional, + Union, + Callable, + Generic, + Iterable, + Iterator, + Protocol, + TypeVar, + Reversible, + List, + Tuple, + Type, +) +from typing_extensions import ( + Concatenate, + ParamSpec, + TypeVarTuple, + Unpack, + TypeAlias, + overload, +) + +_P = ParamSpec("_P") +_In = TypeVar("_In", bound=Iterable) +_Out = TypeVar("_Out", bound=Iterable) +_PipeT = TypeVar("_PipeT", bound="Pipe") + +_T = TypeVar("_T") +_T1 = TypeVar("_T1") +_T2 = TypeVar("_T2") +_T3 = TypeVar("_T3") +_T4 = TypeVar("_T4") +_T5 = TypeVar("_T5") +_S = TypeVar("_S") + +_T_contra = TypeVar("_T_contra", contravariant=True) + +class SupportsRichComparison(Protocol[_T_contra]): + def __lt__(self, other: _T_contra, /) -> bool: ... + def __gt__(self, other: _T_contra, /) -> bool: ... + +_SupportsRichComparisonT = TypeVar( + "_SupportsRichComparisonT", bound=SupportsRichComparison +) + +class Pipe(Generic[_In, _Out, _P]): + args: _P.args + kwargs: _P.kwargs + function: Callable[Concatenate[_In, _P], _Out] + + def __init__( + self, + function: Callable[Concatenate[_In, _P], _Out], + *args: _P.args, + **kwargs: _P.kwargs, + ) -> None: ... + def __ror__(self, other: _In) -> _Out: ... + def __call__(self: _PipeT, *args: _P.args, **kwargs: _P.kwargs) -> _PipeT: ... + def __get__( + self: _PipeT, instance: Optional[_T], owner: Optional[Type[_T]] = ... + ) -> _PipeT: ... + +class _TakePipe(Pipe[Iterable[_T], Iterator[_T], [int]], Generic[_T]): + def __call__(self, qte: int) -> _TakePipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[_T]: ... + +take: _TakePipe + +class _DedupPipe(Pipe[Iterable[_T], Iterator[_T], [Callable[[_T], _T]]], Generic[_T]): + def __call__(self, key: Callable[[_T], _T] = ...) -> _DedupPipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[_T]: ... + +dedup: _DedupPipe + +class _TailPipe(Pipe[Iterable[_T], Sequence[_T], [int]], Generic[_T]): + def __call__(self, qte: int) -> _TailPipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Sequence[_T]: ... + +tail: _TailPipe + +class _SkipPipe(Pipe[Iterable[_T], Iterator[_T], [int]], Generic[_T]): + def __call__(self, qte: int) -> _SkipPipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[_T]: ... + +skip: _SkipPipe + +class _UniqPipe(Pipe[Iterable[_T], Iterator[_T], [Callable[[_T], _T]]], Generic[_T]): + def __call__(self, key: Callable[[_T], _T] = ...) -> _UniqPipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[_T]: ... + +uniq: _UniqPipe + +class _EnumeratePipe(Pipe[Iterable[_T], Iterator[Tuple[int, _T]], [int]], Generic[_T]): + def __call__(self, start: int = ...) -> _EnumeratePipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[Tuple[int, _T]]: ... + +enumerate: _EnumeratePipe + +class _PermutationsPipe( + Pipe[Iterable[_T], Iterable[Tuple[_T, ...]], [int]], Generic[_T] +): + def __call__(self, r: Optional[int] = ...) -> _PermutationsPipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterable[Tuple[_T, ...]]: ... + +permutations: _PermutationsPipe + +class _NetcatPipe(Pipe[Iterable[_T1], Iterator[_T2], [str, int]], Generic[_T1, _T2]): + def __call__(self, host: str, port: int) -> _NetcatPipe[_T1, _T2]: ... + def __ror__(self, iterable: Iterable[_T1]) -> Iterator[_T2]: ... + +netcat: _NetcatPipe + +RIterable: TypeAlias = Iterable[Union[_T, Iterable["RIterable[_T]"]]] + +class _TraversePipe( + Pipe[RIterable[_T], Iterable[_T], []], + Generic[_T], +): + def __call__(self) -> _TraversePipe[_T]: ... + def __ror__(self, iterable: RIterable[_T]) -> Iterable[_T]: ... + +traverse: _TraversePipe + +class _TeePipe(Pipe[Iterable[_T], Iterator[_T], []], Generic[_T]): + def __call__(self) -> _TeePipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[_T]: ... + +tee: _TeePipe + +class _SelectPipe( + Pipe[Iterable[_T1], Iterable[_T2], [Callable[[_T1], _T2]]], Generic[_T1, _T2] +): + def __call__(self, selector: Callable[[_T1], _T2]) -> _SelectPipe[_T1, _T2]: ... + def __ror__(self, iterable: Iterable[_T1]) -> Iterable[_T2]: ... + +select: _SelectPipe +map: _SelectPipe + +class _WherePipe(Pipe[Iterable[_T], Iterator[_T], [Callable[[_T], bool]]], Generic[_T]): + def __call__(self, predicate: Callable[[_T], bool]) -> _WherePipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[_T]: ... + +where: _WherePipe +filter: _WherePipe + +class _TakeWhilePipe( + Pipe[Iterable[_T], Iterable[_T], [Callable[[_T], bool]]], Generic[_T] +): + def __call__(self, predicate: Callable[[_T], bool]) -> _TakeWhilePipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterable[_T]: ... + +take_while: _TakeWhilePipe + +class _SkipWhilePipe( + Pipe[Iterable[_T], Iterable[_T], [Callable[[_T], bool]]], Generic[_T] +): + def __call__(self, predicate: Callable[[_T], bool]) -> _SkipWhilePipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterable[_T]: ... + +skip_while: _SkipWhilePipe + +class _GroupbyPipe( + Pipe[ + Iterable[_T], + Iterator[Tuple[_SupportsRichComparisonT, Iterator[_T]]], + [Callable[[_T], _SupportsRichComparisonT]], + ], + Generic[_T, _SupportsRichComparisonT], +): + def __call__( + self, keyfunc: Callable[[_T], _SupportsRichComparisonT] + ) -> _GroupbyPipe[_T, _SupportsRichComparisonT]: ... + def __ror__( + self, iterable: Iterable[_T] + ) -> Iterator[Tuple[_SupportsRichComparisonT, Iterator[_T]]]: ... + +groupby: _GroupbyPipe + +class _SortPipe( + Pipe[ + Iterable[_T], + List[_SupportsRichComparisonT], + [Callable[[_T], _SupportsRichComparisonT], bool], + ], + Generic[_T, _SupportsRichComparisonT], +): + @overload + def __call__(self, key: None = None, reverse: bool = ...) -> _SortPipe[_T, _T]: ... + @overload + def __call__( + self, + key: Optional[Callable[[_T], _SupportsRichComparisonT]] = ..., + reverse: bool = ..., + ) -> _SortPipe[_T, _SupportsRichComparisonT]: ... + def __ror__(self, iterable: Iterable[_T]) -> List[_SupportsRichComparisonT]: ... + +sort: _SortPipe + +class _ReversePipe(Pipe[Reversible[_T], Iterator[_T], []], Generic[_T]): + def __call__(self) -> _ReversePipe[_T]: ... + def __ror__(self, iterable: Reversible[_T]) -> Iterator[_T]: ... + +reverse: _ReversePipe + +class _TransposePipe(Pipe[Iterable[Iterable[_T]], List[Iterable[_T]], []], Generic[_T]): + def __call__(self) -> _TransposePipe[_T]: ... + def __ror__(self, iterable: Iterable[Iterable[_T]]) -> List[Iterable[_T]]: ... + +transpose: _TransposePipe + +class _BatchedPipe(Pipe[Iterable[_T], Iterator[Tuple[_T, ...]], [int]], Generic[_T]): + def __call__(self, n: int) -> _BatchedPipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[Tuple[_T, ...]]: ... + +batched: _BatchedPipe + +class _ChainPipe(Pipe[Iterable[Iterable[_T]], Iterator[_T], []], Generic[_T]): + def __call__(self) -> _ChainPipe[_T]: ... + def __ror__(self, iterable: Iterable[Iterable[_T]]) -> Iterator[_T]: ... + +chain: _ChainPipe + +_CombinedT = TypeVarTuple("_CombinedT") + +class _ChainWithPipe( + Pipe[Iterable[_T], Iterator[Union[_T, Unpack[_CombinedT]]], _P], + Generic[_T, Unpack[_CombinedT], _P], +): + @overload + def __call__( + self, iterable1: Iterable[_T1], *, strict: bool = ... + ) -> _ChainWithPipe[_T, _T1, [Iterable[_T1], bool]]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + *, + strict: bool = ..., + ) -> _ChainWithPipe[_T, _T1, _T2, [Iterable[_T1], Iterable[_T2], bool]]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + iterable3: Iterable[_T3], + *, + strict: bool = ..., + ) -> _ChainWithPipe[ + _T, _T1, _T2, _T3, [Iterable[_T1], Iterable[_T2], Iterable[_T3], bool] + ]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + iterable3: Iterable[_T3], + iterable4: Iterable[_T4], + *, + strict: bool = ..., + ) -> _ChainWithPipe[ + _T, + _T1, + _T2, + _T3, + _T4, + [Iterable[_T1], Iterable[_T2], Iterable[_T3], Iterable[_T4], bool], + ]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + iterable3: Iterable[_T3], + iterable4: Iterable[_T4], + iterable5: Iterable[_T5], + *, + strict: bool = ..., + ) -> _ChainWithPipe[ + _T, + _T1, + _T2, + _T3, + _T4, + _T5, + [ + Iterable[_T1], + Iterable[_T2], + Iterable[_T3], + Iterable[_T4], + Iterable[_T5], + bool, + ], + ]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + iterable3: Iterable[_T3], + iterable4: Iterable[_T4], + iterable5: Iterable[_T5], + *iterables: Iterable[object], + strict: bool = ..., + ) -> _ChainWithPipe[_T, object, [Iterable[object], bool]]: ... + def __ror__( + self, iterable: Iterable[_T] + ) -> Iterator[Union[_T, Unpack[_CombinedT]]]: ... + +chain_with: _ChainWithPipe + +class _ISlicePipe( + Pipe[Iterable[_T], Iterator[_T], [int, Optional[int], Optional[int]]], Generic[_T] +): + @overload + def __init__(self, stop: Optional[int]) -> None: ... + @overload + def __init__( + self, start: Optional[int], stop: Optional[int], step: Optional[int] = ... + ) -> None: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[_T]: ... + +islice: _ISlicePipe + +_CombinedTupleT = TypeVarTuple("_CombinedTupleT") + +class _IZipPipe( + Pipe[Iterable[_T], Iterator[Tuple[_T, Unpack[_CombinedTupleT]]], _P], + Generic[_T, Unpack[_CombinedTupleT], _P], +): + # @overload + # def __call__(self) -> _IZipPipe[_T, None, []]: ... + @overload + def __call__( + self, iterable1: Iterable[_T1] + ) -> _IZipPipe[_T, _T1, [Iterable[_T1]]]: ... + @overload + def __call__( + self, iterable1: Iterable[_T1], iterable2: Iterable[_T2] + ) -> _IZipPipe[_T, _T1, _T2, [Iterable[_T1], Iterable[_T2]]]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + iterable3: Iterable[_T3], + ) -> _IZipPipe[ + _T, _T1, _T2, _T3, [Iterable[_T1], Iterable[_T2], Iterable[_T3]] + ]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + iterable3: Iterable[_T3], + iterable4: Iterable[_T4], + ) -> _IZipPipe[ + _T, + _T1, + _T2, + _T3, + _T4, + [Iterable[_T1], Iterable[_T2], Iterable[_T3], Iterable[_T4]], + ]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + iterable3: Iterable[_T3], + iterable4: Iterable[_T4], + iterable5: Iterable[_T5], + ) -> _IZipPipe[ + _T, + _T1, + _T2, + _T3, + _T4, + _T5, + [Iterable[_T1], Iterable[_T2], Iterable[_T3], Iterable[_T4], Iterable[_T5]], + ]: ... + @overload + def __call__( + self, + iterable1: Iterable[_T1], + iterable2: Iterable[_T2], + iterable3: Iterable[_T3], + iterable4: Iterable[_T4], + iterable5: Iterable[_T5], + *iterables: Iterable[object], + ) -> _IZipPipe[_T, object, [Iterable[object]]]: ... + def __ror__( + self, iterable: Iterable[_T] + ) -> Iterator[Tuple[_T, Unpack[_CombinedTupleT]]]: ... + +izip: _IZipPipe + +class _TPipe(Pipe[Iterable[_T], Iterator[_T], [_T]], Generic[_T]): + def __call__(self, y: _T) -> _TPipe[_T]: ... + def __ror__(self, iterable: Iterable[_T]) -> Iterator[_T]: ... + +t: _TPipe + +class _CastPipe(Pipe[Iterable[_T], _S, [Type[_S]]], Generic[_T, _S]): + def __call__(self, type_: Type[_S]) -> _CastPipe[_T, _S]: ... + def __ror__(self, iterable: Iterable[_T]) -> _S: ... + +cast: _CastPipe