from typing import Any
from typing import Callable
from typing import Mapping
from typing import Sequence


def _map_batch(batch: Any, map_fn: Callable) -> Any:
    """Recursively maps a function to all items in a batch.

    Args:
        batch: Nested lists and dictionaries.
        map_fn: A function to invoke on each element.

    Returns:
        Collections: The result of applying ``map_fn`` on each element of the ``batch``.
        The type of ``batch`` is preserved.
    """
    if isinstance(batch, Mapping):
        return {k: _map_batch(v, map_fn) for k, v in batch.items()}

    if isinstance(batch, Sequence) and not isinstance(batch, (str, bytes)):
        try:
            return type(batch)(_map_batch(x, map_fn) for x in batch)  # type: ignore
        except TypeError:
            return [_map_batch(x, map_fn) for x in batch]
    return map_fn(batch)