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

from remedapy.decorator import make_data_last

Key = TypeVar('Key', str, int)
T = TypeVar('T')


@overload
def pick(data: dict[Key, T], keys: Collection[Key], /) -> dict[Key, T]: ...


@overload
def pick(keys: Collection[Key], /) -> Callable[[dict[Key, T]], dict[Key, T]]: ...


@make_data_last
def pick(data: dict[Key, T], keys: Collection[Key], /) -> dict[Key, T]:
    """
    Creates a new dict by picking the given keys from the given dict.

    Parameters
    ----------
    data: dict[K, V]
        The dict to pick keys from.
    keys: Collection[K]
        The keys to pick from the dict.

    Returns
    -------
    dict[K, V]
        The new dict with the picked keys.

    See Also
    --------
    omit
    pick_by

    Examples
    --------
    Data first:
    >>> R.pick({'a': 1, 'b': 2, 'c': 3, 'd': 4}, ['a', 'd'])
    {'a': 1, 'd': 4}

    Data last:
    >>> R.pipe({'a': 1, 'b': 2, 'c': 3, 'd': 4}, R.pick(['a', 'd']))
    {'a': 1, 'd': 4}

    """
    return {k: v for k, v in data.items() if k in keys}
