Source code for pythonwrench.collections.reducers

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import operator
from typing import (
    Any,
    Callable,
    Iterable,
    Iterator,
    Optional,
    Tuple,
    Type,
    TypeVar,
    overload,
)

from pythonwrench.functools import function_alias
from pythonwrench.typing.checks import isinstance_generic
from pythonwrench.typing.classes import (
    SupportsAdd,
    SupportsAnd,
    SupportsMatmul,
    SupportsMul,
    SupportsOr,
)

T = TypeVar("T")
T_SupportsAdd = TypeVar("T_SupportsAdd", bound=SupportsAdd)
T_SupportsAnd = TypeVar("T_SupportsAnd", bound=SupportsAnd)
T_SupportsMul = TypeVar("T_SupportsMul", bound=SupportsMul)
T_SupportsOr = TypeVar("T_SupportsOr", bound=SupportsOr)
T_SupportsMatmul = TypeVar("T_SupportsMatmul", bound=SupportsMatmul)


@overload
def reduce_add(
    args: Iterable[T_SupportsAdd],
    /,
    *,
    start: T_SupportsAdd,
) -> T_SupportsAdd: ...


@overload
def reduce_add(
    *args: T_SupportsAdd,
    start: T_SupportsAdd,
) -> T_SupportsAdd: ...


@overload
def reduce_add(
    arg0: T_SupportsAdd,
    /,
    *args: T_SupportsAdd,
    start: Optional[T_SupportsAdd] = None,
) -> T_SupportsAdd: ...



[docs]
def reduce_add(*args, start=None):
    """Reduce elements using "add" operator (+)."""
    return _reduce(*args, start=start, op_fn=operator.add, type_=SupportsAdd)



@overload
def reduce_and(
    args: Iterable[T_SupportsAnd],
    /,
    *,
    start: T_SupportsAnd,
) -> T_SupportsAnd: ...


@overload
def reduce_and(
    *args: T_SupportsAnd,
    start: T_SupportsAnd,
) -> T_SupportsAnd: ...


@overload
def reduce_and(
    arg0: T_SupportsAnd,
    /,
    *args: T_SupportsAnd,
    start: Optional[T_SupportsAnd] = None,
) -> T_SupportsAnd: ...



[docs]
def reduce_and(*args, start=None):
    """Reduce elements using "and" operator (&)."""
    return _reduce(*args, start=start, op_fn=operator.and_, type_=SupportsAnd)



@overload
def reduce_matmul(
    args: Iterable[T_SupportsMatmul],
    /,
    *,
    start: T_SupportsMatmul,
) -> T_SupportsMatmul: ...


@overload
def reduce_matmul(
    *args: T_SupportsMatmul,
    start: T_SupportsMatmul,
) -> T_SupportsMatmul: ...


@overload
def reduce_matmul(
    arg0: T_SupportsMatmul,
    /,
    *args: T_SupportsMatmul,
    start: Optional[T_SupportsMatmul] = None,
) -> T_SupportsMatmul: ...



[docs]
def reduce_matmul(*args, start=None):
    """Reduce elements using "mul" operator (*)."""
    return _reduce(*args, start=start, op_fn=operator.matmul, type_=SupportsMatmul)



@overload
def reduce_mul(
    args: Iterable[T_SupportsMul],
    /,
    *,
    start: T_SupportsMul,
) -> T_SupportsMul: ...


@overload
def reduce_mul(
    *args: T_SupportsMul,
    start: T_SupportsMul,
) -> T_SupportsMul: ...


@overload
def reduce_mul(
    arg0: T_SupportsMul,
    /,
    *args: T_SupportsMul,
    start: Optional[T_SupportsMul] = None,
) -> T_SupportsMul: ...



[docs]
def reduce_mul(*args, start=None):
    """Reduce elements using "mul" operator (*)."""
    return _reduce(*args, start=start, op_fn=operator.mul, type_=SupportsMul)



@overload
def reduce_or(
    args: Iterable[T_SupportsOr],
    /,
    *,
    start: T_SupportsOr,
) -> T_SupportsOr: ...


@overload
def reduce_or(
    *args: T_SupportsOr,
    start: T_SupportsOr,
) -> T_SupportsOr: ...


@overload
def reduce_or(
    arg0: T_SupportsOr,
    /,
    *args: T_SupportsOr,
    start: Optional[T_SupportsOr] = None,
) -> T_SupportsOr: ...



[docs]
def reduce_or(*args, start=None):
    """Reduce elements using "or" operator (|)."""
    return _reduce(*args, start=start, op_fn=operator.or_, type_=SupportsOr)



def _reduce(
    *args,
    start: Optional[T] = None,
    op_fn: Callable[[T, T], T],
    type_: Type[T],
) -> T:
    if isinstance_generic(args, Tuple[Iterable[type_]]):
        it_or_args = args[0]
    elif isinstance_generic(args, Tuple[type_, ...]):
        it_or_args = args
    else:
        msg = f"Invalid positional arguments {args}. (expected {Tuple[type_, ...]} or {Tuple[Iterable[type_]]})"
        raise TypeError(msg)

    it: Iterator[T] = iter(it_or_args)

    if isinstance(start, type_):
        accumulator = start
    elif start is None or start is ...:
        try:
            accumulator = next(it)
        except StopIteration:
            msg = f"Invalid combinaison of arguments {args=} and {start=}. (expected at least 1 non-empty argument or start object that supports operator.)"
            raise ValueError(msg)
    else:
        msg = f"Invalid argument type {type(start)}."
        raise TypeError(msg)

    for arg in it:
        accumulator = op_fn(accumulator, arg)
    return accumulator


@overload
def sum(
    args: Iterable[T_SupportsAdd],
    /,
    *,
    start: T_SupportsAdd = 0,
) -> T_SupportsAdd: ...


@overload
def sum(
    *args: T_SupportsAdd,
    start: T_SupportsAdd = 0,
) -> T_SupportsAdd: ...


@overload
def sum(
    arg0: T_SupportsAdd,
    /,
    *args: T_SupportsAdd,
    start: Optional[T_SupportsAdd] = 0,
) -> T_SupportsAdd: ...



[docs]
def sum(*args, start: Any = 0):
    """Compute sum of elements."""
    return reduce_add(*args, start=start)



@overload
def prod(
    args: Iterable[T_SupportsMul],
    /,
    *,
    start: T_SupportsMul = 1,
) -> T_SupportsMul: ...


@overload
def prod(
    *args: T_SupportsMul,
    start: T_SupportsMul = 1,
) -> T_SupportsMul: ...


@overload
def prod(
    arg0: T_SupportsMul,
    /,
    *args: T_SupportsMul,
    start: Optional[T_SupportsMul] = 1,
) -> T_SupportsMul: ...



[docs]
def prod(*args, start: Any = 1):
    """Compute product of elements."""
    return reduce_mul(*args, start=start)




[docs]
@function_alias(reduce_and)
def intersect(*args, **kwargs): ...




[docs]
@function_alias(reduce_or)
def union(*args, **kwargs): ...