pyochain 0.5.0__py3-none-any.whl

pyochain/_iter/_groups.py

@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Iterator
+from functools import partial
+from typing import TYPE_CHECKING, Any, overload
+
+import cytoolz as cz
+import more_itertools as mit
+
+from .._core import IterWrapper
+
+if TYPE_CHECKING:
+    from .._dict import Dict
+    from ._main import Iter
+
+
+class BaseGroups[T](IterWrapper[T]):
+    def reduce_by[K](
+        self, key: Callable[[T], K], binop: Callable[[T, T], T]
+    ) -> Dict[K, T]:
+        """
+        Perform a simultaneous groupby and reduction.
+
+        Args:
+            key: Function to compute the key for grouping.
+            binop: Binary operation to reduce the grouped elements.
+        Example:
+        ```python
+        >>> from collections.abc import Iterable
+        >>> import pyochain as pc
+        >>> from operator import add, mul
+        >>>
+        >>> def is_even(x: int) -> bool:
+        ...     return x % 2 == 0
+        >>>
+        >>> def group_reduce(data: Iterable[int]) -> int:
+        ...     return pc.Iter.from_(data).reduce(add)
+        >>>
+        >>> data = pc.Seq([1, 2, 3, 4, 5])
+        >>> data.iter().reduce_by(is_even, add).unwrap()
+        {False: 9, True: 6}
+        >>> data.iter().group_by(is_even).map_values(group_reduce).unwrap()
+        {False: 9, True: 6}
+
+        ```
+        But the former does not build the intermediate groups, allowing it to operate in much less space.
+
+        This makes it suitable for larger datasets that do not fit comfortably in memory
+
+        Simple Examples:
+        ```python
+        >>> pc.Iter.from_([1, 2, 3, 4, 5]).reduce_by(is_even, add).unwrap()
+        {False: 9, True: 6}
+        >>> pc.Iter.from_([1, 2, 3, 4, 5]).reduce_by(is_even, mul).unwrap()
+        {False: 15, True: 8}
+
+        ```
+        """
+        from .._dict import Dict
+
+        return Dict(self.into(partial(cz.itertoolz.reduceby, key, binop)))
+
+    def group_by[K](self, on: Callable[[T], K]) -> Dict[K, list[T]]:
+        """
+        Group elements by key function and return a Dict result.
+
+        Args:
+            on: Function to compute the key for grouping.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> names = [
+        ...     "Alice",
+        ...     "Bob",
+        ...     "Charlie",
+        ...     "Dan",
+        ...     "Edith",
+        ...     "Frank",
+        ... ]
+        >>> pc.Iter.from_(names).group_by(len).sort()
+        ... # doctest: +NORMALIZE_WHITESPACE
+        Dict({
+            3: ['Bob', 'Dan'],
+            5: ['Alice', 'Edith', 'Frank'],
+            7: ['Charlie']
+        })
+        >>>
+        >>> iseven = lambda x: x % 2 == 0
+        >>> pc.Iter.from_([1, 2, 3, 4, 5, 6, 7, 8]).group_by(iseven)
+        ... # doctest: +NORMALIZE_WHITESPACE
+        Dict({
+            False: [1, 3, 5, 7],
+            True: [2, 4, 6, 8]
+        })
+
+        ```
+        Non-callable keys imply grouping on a member.
+        ```python
+        >>> data = [
+        ...     {"name": "Alice", "gender": "F"},
+        ...     {"name": "Bob", "gender": "M"},
+        ...     {"name": "Charlie", "gender": "M"},
+        ... ]
+        >>> pc.Iter.from_(data).group_by("gender").sort()
+        ... # doctest: +NORMALIZE_WHITESPACE
+        Dict({
+            'F': [
+                {'name': 'Alice', 'gender': 'F'}
+            ],
+            'M': [
+                {'name': 'Bob', 'gender': 'M'},
+                {'name': 'Charlie', 'gender': 'M'}
+            ]
+        })
+
+        ```
+        """
+        from .._dict import Dict
+
+        return Dict(self.into(partial(cz.itertoolz.groupby, on)))
+
+    def frequencies(self) -> Dict[T, int]:
+        """
+        Find number of occurrences of each value in the iterable.
+        ```python
+        >>> import pyochain as pc
+        >>> data = ["cat", "cat", "ox", "pig", "pig", "cat"]
+        >>> pc.Iter.from_(data).frequencies().unwrap()
+        {'cat': 3, 'ox': 1, 'pig': 2}
+
+        ```
+        """
+        from .._dict import Dict
+
+        return Dict(self.into(cz.itertoolz.frequencies))
+
+    def count_by[K](self, key: Callable[[T], K]) -> Dict[K, int]:
+        """
+        Count elements of a collection by a key function.
+
+        Args:
+            key: Function to compute the key for counting.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> pc.Iter.from_(["cat", "mouse", "dog"]).count_by(len).unwrap()
+        {3: 2, 5: 1}
+        >>> def iseven(x):
+        ...     return x % 2 == 0
+        >>> pc.Iter.from_([1, 2, 3]).count_by(iseven).unwrap()
+        {False: 2, True: 1}
+
+        ```
+        """
+        from .._dict import Dict
+
+        return Dict(self.into(partial(cz.recipes.countby, key)))
+
+    @overload
+    def group_by_transform(
+        self,
+        keyfunc: None = None,
+        valuefunc: None = None,
+        reducefunc: None = None,
+    ) -> Iter[tuple[T, Iterator[T]]]: ...
+    @overload
+    def group_by_transform[U](
+        self,
+        keyfunc: Callable[[T], U],
+        valuefunc: None,
+        reducefunc: None,
+    ) -> Iter[tuple[U, Iterator[T]]]: ...
+    @overload
+    def group_by_transform[V](
+        self,
+        keyfunc: None,
+        valuefunc: Callable[[T], V],
+        reducefunc: None,
+    ) -> Iter[tuple[T, Iterator[V]]]: ...
+    @overload
+    def group_by_transform[U, V](
+        self,
+        keyfunc: Callable[[T], U],
+        valuefunc: Callable[[T], V],
+        reducefunc: None,
+    ) -> Iter[tuple[U, Iterator[V]]]: ...
+    @overload
+    def group_by_transform[W](
+        self,
+        keyfunc: None,
+        valuefunc: None,
+        reducefunc: Callable[[Iterator[T]], W],
+    ) -> Iter[tuple[T, W]]: ...
+    @overload
+    def group_by_transform[U, W](
+        self,
+        keyfunc: Callable[[T], U],
+        valuefunc: None,
+        reducefunc: Callable[[Iterator[T]], W],
+    ) -> Iter[tuple[U, W]]: ...
+    @overload
+    def group_by_transform[V, W](
+        self,
+        keyfunc: None,
+        valuefunc: Callable[[T], V],
+        reducefunc: Callable[[Iterator[V]], W],
+    ) -> Iter[tuple[T, W]]: ...
+    @overload
+    def group_by_transform[U, V, W](
+        self,
+        keyfunc: Callable[[T], U],
+        valuefunc: Callable[[T], V],
+        reducefunc: Callable[[Iterator[V]], W],
+    ) -> Iter[tuple[U, W]]: ...
+    def group_by_transform[U, V](
+        self,
+        keyfunc: Callable[[T], U] | None = None,
+        valuefunc: Callable[[T], V] | None = None,
+        reducefunc: Any = None,
+    ) -> Iter[tuple[Any, ...]]:
+        """
+        An extension of itertools.groupby that can apply transformations to the grouped data.
+
+        Args:
+            keyfunc: Function to compute the key for grouping. Defaults to None.
+            valuefunc: Function to transform individual items after grouping. Defaults to None.
+            reducefunc: Function to transform each group of items. Defaults to None.
+
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> data = pc.Iter.from_("aAAbBBcCC")
+        >>> data.group_by_transform(
+        ...     lambda k: k.upper(), lambda v: v.lower(), lambda g: "".join(g)
+        ... ).into(list)
+        [('A', 'aaa'), ('B', 'bbb'), ('C', 'ccc')]
+
+        ```
+        Each optional argument defaults to an identity function if not specified.
+
+        group_by_transform is useful when grouping elements of an iterable using a separate iterable as the key.
+
+        To do this, zip the iterables and pass a keyfunc that extracts the first element and a valuefunc that extracts the second element:
+
+        Note that the order of items in the iterable is significant.
+
+        Only adjacent items are grouped together, so if you don't want any duplicate groups, you should sort the iterable by the key function.
+
+        Example:
+        ```python
+        >>> from operator import itemgetter
+        >>> data = pc.Iter.from_([0, 0, 1, 1, 1, 2, 2, 2, 3])
+        >>> data.zip("abcdefghi").group_by_transform(itemgetter(0), itemgetter(1)).map(
+        ...     lambda kv: (kv[0], "".join(kv[1]))
+        ... ).into(list)
+        [(0, 'ab'), (1, 'cde'), (2, 'fgh'), (3, 'i')]
+
+        ```
+        """
+
+        def _group_by_transform(data: Iterable[T]) -> Iterator[tuple[Any, ...]]:
+            return mit.groupby_transform(data, keyfunc, valuefunc, reducefunc)
+
+        return self.apply(_group_by_transform)

pyochain/_iter/_joins.py

@@ -0,0 +1,407 @@
+from __future__ import annotations
+
+import itertools
+from collections.abc import Callable, Generator, Iterable, Iterator
+from typing import TYPE_CHECKING, Any, overload
+
+import cytoolz as cz
+import more_itertools as mit
+
+from .._core import IterWrapper
+
+if TYPE_CHECKING:
+    from ._main import Iter
+
+
+class BaseJoins[T](IterWrapper[T]):
+    @overload
+    def zip[T1](
+        self, iter1: Iterable[T1], /, *, strict: bool = ...
+    ) -> Iter[tuple[T, T1]]: ...
+    @overload
+    def zip[T1, T2](
+        self,
+        iter1: Iterable[T1],
+        iter2: Iterable[T2],
+        /,
+        *,
+        strict: bool = ...,
+    ) -> Iter[tuple[T, T1, T2]]: ...
+    @overload
+    def zip[T1, T2, T3](
+        self,
+        iter1: Iterable[T1],
+        iter2: Iterable[T2],
+        iter3: Iterable[T3],
+        /,
+        *,
+        strict: bool = ...,
+    ) -> Iter[tuple[T, T1, T2, T3]]: ...
+    @overload
+    def zip[T1, T2, T3, T4](
+        self,
+        iter1: Iterable[T1],
+        iter2: Iterable[T2],
+        iter3: Iterable[T3],
+        iter4: Iterable[T4],
+        /,
+        *,
+        strict: bool = ...,
+    ) -> Iter[tuple[T, T1, T2, T3, T4]]: ...
+    def zip(
+        self, *others: Iterable[Any], strict: bool = False
+    ) -> Iter[tuple[Any, ...]]:
+        """
+        Zip with other iterables, optionally strict.
+
+        Args:
+            *others: Other iterables to zip with.
+            strict: Whether to enforce equal lengths of iterables. Defaults to False.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> pc.Iter.from_([1, 2]).zip([10, 20]).into(list)
+        [(1, 10), (2, 20)]
+        >>> pc.Iter.from_(["a", "b"]).zip([1, 2, 3]).into(list)
+        [('a', 1), ('b', 2)]
+
+        ```
+        """
+        return self.apply(zip, *others, strict=strict)
+
+    def zip_offset[U](
+        self,
+        *others: Iterable[T],
+        offsets: list[int],
+        longest: bool = False,
+        fillvalue: U = None,
+    ) -> Iter[tuple[T | U, ...]]:
+        """
+        Zip the input iterables together, but offset the i-th iterable by the i-th item in offsets.
+
+        Args:
+            *others: Other iterables to zip with.
+            offsets: List of integers specifying the offsets for each iterable.
+            longest: Whether to continue until the longest iterable is exhausted. Defaults to False.
+            fillvalue: Value to use for missing elements. Defaults to None.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> data = pc.Seq("0123")
+        >>> data.iter().zip_offset("abcdef", offsets=(0, 1)).into(list)
+        [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e')]
+
+        ```
+        This can be used as a lightweight alternative to SciPy or pandas to analyze data sets in which some series have a lead or lag relationship.
+
+        By default, the sequence will end when the shortest iterable is exhausted.
+
+        To continue until the longest iterable is exhausted, set longest to True.
+        ```python
+        >>> data.iter().zip_offset("abcdef", offsets=(0, 1), longest=True).into(list)
+        [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e'), (None, 'f')]
+
+        ```
+        """
+
+        def _zip_offset(data: Iterable[T]) -> Iterator[tuple[T | U, ...]]:
+            return mit.zip_offset(
+                data,
+                *others,
+                offsets=offsets,
+                longest=longest,
+                fillvalue=fillvalue,
+            )
+
+        return self.apply(_zip_offset)
+
+    @overload
+    def zip_broadcast[T1](
+        self,
+        iter1: Iterable[T1],
+        /,
+        *,
+        strict: bool = False,
+    ) -> Iter[tuple[T, T1]]: ...
+    @overload
+    def zip_broadcast[T1, T2](
+        self,
+        iter1: Iterable[T1],
+        iter2: Iterable[T2],
+        /,
+        *,
+        strict: bool = False,
+    ) -> Iter[tuple[T, T1, T2]]: ...
+    @overload
+    def zip_broadcast[T1, T2, T3](
+        self,
+        iter1: Iterable[T1],
+        iter2: Iterable[T2],
+        iter3: Iterable[T3],
+        /,
+        *,
+        strict: bool = False,
+    ) -> Iter[tuple[T, T1, T2, T3]]: ...
+    @overload
+    def zip_broadcast[T1, T2, T3, T4](
+        self,
+        iter1: Iterable[T1],
+        iter2: Iterable[T2],
+        iter3: Iterable[T3],
+        iter4: Iterable[T4],
+        /,
+        *,
+        strict: bool = False,
+    ) -> Iter[tuple[T, T1, T2, T3, T4]]: ...
+    def zip_broadcast(
+        self, *others: Iterable[Any], strict: bool = False
+    ) -> Iter[tuple[Any, ...]]:
+        """
+        Version of zip that "broadcasts" any scalar (i.e., non-iterable) items into output tuples.
+
+        `str` and `bytes` are not treated as iterables.
+
+        If the strict keyword argument is True, then UnequalIterablesError will be raised if any of the iterables have different lengths.
+        Args:
+            *others: Other iterables or scalars to zip with.
+            strict: Whether to enforce equal lengths of iterables. Defaults to False.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> data = pc.Iter.from_([1, 2, 3])
+        >>> other = ["a", "b", "c"]
+        >>> scalar = "_"
+        >>> data.zip_broadcast(other, scalar).into(list)
+        [(1, 'a', '_'), (2, 'b', '_'), (3, 'c', '_')]
+
+        ```
+        """
+
+        def _zip_broadcast(
+            *objects: Iterable[Any],
+        ) -> Generator[tuple[Iterable[Any], ...] | tuple[object, ...], Any, None]:
+            """from more_itertools.zip_broadcast"""
+
+            def is_scalar(obj: Any) -> bool:
+                if isinstance(obj, (str, bytes)):
+                    return True
+                try:
+                    iter(obj)
+                except TypeError:
+                    return True
+                else:
+                    return False
+
+            size = len(objects)
+            if not size:
+                return
+
+            new_item: list[object] = [None] * size
+            iterables: list[Iterator[Any]] = []
+            iterable_positions: list[int] = []
+            for i, obj in enumerate(objects):
+                if is_scalar(obj):
+                    new_item[i] = obj
+                else:
+                    iterables.append(iter(obj))
+                    iterable_positions.append(i)
+
+            if not iterables:
+                yield tuple(objects)
+                return
+
+            zipper = mit.zip_equal if strict else zip
+            for item in zipper(*iterables):
+                for i, new_item[i] in zip(iterable_positions, item):
+                    pass
+                yield tuple(new_item)
+
+        return self.apply(_zip_broadcast, *others)
+
+    @overload
+    def zip_equal(self) -> Iter[tuple[T]]: ...
+    @overload
+    def zip_equal[T2](self, __iter2: Iterable[T2]) -> Iter[tuple[T, T2]]: ...
+    @overload
+    def zip_equal[T2, T3](
+        self, __iter2: Iterable[T2], __iter3: Iterable[T3]
+    ) -> Iter[tuple[T, T2, T3]]: ...
+    @overload
+    def zip_equal[T2, T3, T4](
+        self,
+        __iter2: Iterable[T2],
+        __iter3: Iterable[T3],
+        __iter4: Iterable[T4],
+    ) -> Iter[tuple[T, T2, T3, T4]]: ...
+    @overload
+    def zip_equal[T2, T3, T4, T5](
+        self,
+        __iter2: Iterable[T2],
+        __iter3: Iterable[T3],
+        __iter4: Iterable[T4],
+        __iter5: Iterable[T5],
+    ) -> Iter[tuple[T, T2, T3, T4, T5]]: ...
+    def zip_equal(self, *others: Iterable[Any]) -> Iter[tuple[Any, ...]]:
+        """
+        `zip` the input *iterables* together but raise `UnequalIterablesError` if they aren't all the same length.
+
+        Args:
+            *others: Other iterables to zip with.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> pc.Iter.from_(range(3)).zip_equal("abc").into(list)
+        [(0, 'a'), (1, 'b'), (2, 'c')]
+        >>> pc.Iter.from_(range(3)).zip_equal("abcd").into(list)
+        ... # doctest: +IGNORE_EXCEPTION_DETAIL
+        Traceback (most recent call last):
+        ...
+        more_itertools.more.UnequalIterablesError: Iterables have different
+        lengths
+
+        ```
+        """
+
+        def _zip_equal(data: Iterable[T]) -> Iterator[tuple[Any, ...]]:
+            return mit.zip_equal(data, *others)
+
+        return self.apply(_zip_equal)
+
+    def zip_longest[U](
+        self, *others: Iterable[T], fill_value: U = None
+    ) -> Iter[tuple[U | T, ...]]:
+        """
+        Zip with other iterables, filling missing values.
+
+        Args:
+            *others: Other iterables to zip with.
+            fill_value: Value to use for missing elements. Defaults to None.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> pc.Iter.from_([1, 2]).zip_longest([10], fill_value=0).into(list)
+        [(1, 10), (2, 0)]
+
+        ```
+        """
+        return self.apply(itertools.zip_longest, *others, fillvalue=fill_value)
+
+    @overload
+    def product(self) -> Iter[tuple[T]]: ...
+    @overload
+    def product[T1](self, iter1: Iterable[T1], /) -> Iter[tuple[T, T1]]: ...
+    @overload
+    def product[T1, T2](
+        self, iter1: Iterable[T1], iter2: Iterable[T2], /
+    ) -> Iter[tuple[T, T1, T2]]: ...
+    @overload
+    def product[T1, T2, T3](
+        self, iter1: Iterable[T1], iter2: Iterable[T2], iter3: Iterable[T3], /
+    ) -> Iter[tuple[T, T1, T2, T3]]: ...
+    @overload
+    def product[T1, T2, T3, T4](
+        self,
+        iter1: Iterable[T1],
+        iter2: Iterable[T2],
+        iter3: Iterable[T3],
+        iter4: Iterable[T4],
+        /,
+    ) -> Iter[tuple[T, T1, T2, T3, T4]]: ...
+
+    def product(self, *others: Iterable[Any]) -> Iter[tuple[Any, ...]]:
+        """
+        Computes the Cartesian product with another iterable.
+        This is the declarative equivalent of nested for-loops.
+
+        It pairs every element from the source iterable with every element from the
+        other iterable.
+
+        Args:
+            *others: Other iterables to compute the Cartesian product with.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> colors = pc.Iter.from_(["blue", "red"])
+        >>> sizes = ["S", "M"]
+        >>> colors.product(sizes).into(list)
+        [('blue', 'S'), ('blue', 'M'), ('red', 'S'), ('red', 'M')]
+
+        ```
+        """
+        return self.apply(itertools.product, *others)
+
+    def diff_at(
+        self,
+        *others: Iterable[T],
+        default: T | None = None,
+        key: Callable[[T], Any] | None = None,
+    ) -> Iter[tuple[T, ...]]:
+        """
+        Return those items that differ between iterables.
+        Each output item is a tuple where the i-th element is from the i-th input iterable.
+
+        If an input iterable is exhausted before others, then the corresponding output items will be filled with *default*.
+        Args:
+            *others: Other iterables to compare with.
+            default: Value to use for missing elements. Defaults to None.
+            key: Function to apply to each item for comparison. Defaults to None.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> data = pc.Seq([1, 2, 3])
+        >>> data.iter().diff_at([1, 2, 10, 100], default=None).into(list)
+        [(3, 10), (None, 100)]
+        >>> data.iter().diff_at([1, 2, 10, 100, 2, 6, 7], default=0).into(list)
+        [(3, 10), (0, 100), (0, 2), (0, 6), (0, 7)]
+
+        A key function may also be applied to each item to use during comparisons:
+        ```python
+        >>> import pyochain as pc
+        >>> pc.Iter.from_(["apples", "bananas"]).diff_at(
+        ...     ["Apples", "Oranges"], key=str.lower
+        ... ).into(list)
+        [('bananas', 'Oranges')]
+
+        ```
+        """
+        return self.apply(cz.itertoolz.diff, *others, default=default, key=key)
+
+    def join[R, K](
+        self,
+        other: Iterable[R],
+        left_on: Callable[[T], K],
+        right_on: Callable[[R], K],
+        left_default: T | None = None,
+        right_default: R | None = None,
+    ) -> Iter[tuple[T, R]]:
+        """
+        Perform a relational join with another iterable.
+
+        Args:
+            other: Iterable to join with.
+            left_on: Function to extract the join key from the left iterable.
+            right_on: Function to extract the join key from the right iterable.
+            left_default: Default value for missing elements in the left iterable. Defaults to None.
+            right_default: Default value for missing elements in the right iterable. Defaults to None.
+        Example:
+        ```python
+        >>> import pyochain as pc
+        >>> colors = pc.Iter.from_(["blue", "red"])
+        >>> sizes = ["S", "M"]
+        >>> colors.join(sizes, left_on=lambda c: c, right_on=lambda s: s).into(list)
+        [(None, 'S'), (None, 'M'), ('blue', None), ('red', None)]
+
+        ```
+        """
+
+        def _join(data: Iterable[T]) -> Iterator[tuple[T, R]]:
+            return cz.itertoolz.join(
+                leftkey=left_on,
+                leftseq=data,
+                rightkey=right_on,
+                rightseq=other,
+                left_default=left_default,
+                right_default=right_default,
+            )
+
+        return self.apply(_join)