from typing import Callable


def map(f: Callable, collection):
    """Transform each element of a collection.

    Examples
    --------
    .. doctest::

        >>> a = ['The', 'quick', 'brown', 'fox']

        >>> hl.eval_expr(hl.map(lambda x: hl.len(x), a))
        [3, 5, 5, 3]

    Parameters
    ----------
    f : function ( (arg) -> :class:`.Expression`)
        Function to transform each element of the collection.
    collection : :class:`.ArrayExpression` or :class:`.SetExpression`
        Collection expression.

    Returns
    -------
    :class:`.ArrayExpression` or :class:`SetExpression`.
        Collection where each element has been transformed by `f`.
    """
    return collection._bin_lambda_method("map", f,
                                         collection.dtype.element_type,
                                         lambda t: collection.dtype.__class__(t))