from collections.abc import Callable, Iterable
from typing import TypeVar, overload

from remedapy.decorator import make_data_last

T = TypeVar('T')


@overload
def chunk(data: Iterable[T], size: int, /) -> Iterable[list[T]]: ...


@overload
def chunk(size: int, /) -> Callable[[Iterable[T]], Iterable[list[T]]]: ...


@make_data_last
def chunk(iterable: Iterable[T], size: int, /) -> Iterable[list[T]]:
    """
    Yields lists of specified size from the iterable.

    Parameters
    ----------
    iterable : Iterable[T]
        Input iterable (positional-only).
    size : int
        Size (length) of the chunks (positional-only).

    Returns
    -------
    Iterable[list[T]]
        Iterable of lists of specified size.

    Examples
    --------
    Data first:
    >>> list(R.chunk([1,2,3,4,5,6], 3))
    [[1, 2, 3], [4, 5, 6]]
    >>> list(R.chunk(range(1, 8), 3))
    [[1, 2, 3], [4, 5, 6], [7]]

    Data last:
    >>> R.pipe(range(1, 8), R.chunk(3), list)
    [[1, 2, 3], [4, 5, 6], [7]]
    >>> R.pipe([1,2,3], R.chunk(2), list)
    [[1, 2], [3]]

    """
    acc: list[T] = []
    for i in iterable:
        acc.append(i)
        if len(acc) == size:
            yield acc
            acc = []
    if acc:
        yield acc
