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

from .decorator import make_data_last

T = TypeVar('T')


@overload
def take(it: Iterable[T], n: int, /) -> Iterable[T]: ...


@overload
def take(n: int, /) -> Callable[[Iterable[T]], Iterable[T]]: ...


@make_data_last
def take(iterable: Iterable[T], n: int, /) -> Iterable[T]:
    """
    Yields the first n elements of the iterable.

    Parameters
    ----------
    iterable : Iterable[T]
        Input iterable (positional-only).
    n : int
        Number of elements to yield (positional-only).

    Returns
    -------
    Iterable[T]
        First n elements of the iterable.

    Examples
    --------
    Data first:
    >>> list(R.take(range(100), 2))
    [0, 1]
    >>> list(R.take([2, 1, 3, 7, 6, 6, 6], 4))
    [2, 1, 3, 7]

    Data last:
    >>> list(R.take([2, 1, 3, 7, 6, 6, 6], 4))
    [2, 1, 3, 7]
    >>> list(R.take(range(100), 2))
    [0, 1]

    """
    if n <= 0:
        return
    for i, x in enumerate(iterable):
        yield x
        if i + 1 == n:
            break
