[Python] Structure-preserving nested map in Python

I'm writing a Python function inspired by map that takes a function f: Callable[[A], B] and an arbitrarily nested structure of iterables (e.g. lists of tuples of lists of ... of objects of type A) and returns an identically nested structure with f being applied to all "leaf nodes" of the original structure. Ideally, it would support dict[Any, A] too and "homogenous" NamedTuples by applying f only to the values and leaving the keys untouched (NamedTuples' fields should all have type A for this to work).

For example:

>>> f = lambda x: x * 2
>>> input = [([1, 2], [3, 4]), {"foo": 5, "bar": 6}]
>>> nested_map(f, input)
[([2, 4], [6, 8]), {"foo": 10, "bar": 12}]


Ideally, it would be type-hint friendly in the sense that it would tell a static checker that the output object has exactly the same structure as the input. Although this is probably way harder to achieve.

I have a kind of working version for this but:

• the code feels really ugly;
• I was unable to properly type-annotate "homogenous" NamedTuples; and
• it doesn't work well with static type checking.

import inspect
from collections.abc import Iterable
from typing import Any, Callable, Generator, NamedTuple, cast

type MappableIterable[A] = list[A] | tuple[A] | dict[Any, A] | Generator[A]

type NestedIterable[A] = A | MappableIterable["NestedIterable[A]"]

def nested_map[A, B](
f: Callable[[A], B], nested_structure: NestedIterable[A]
) -> NestedIterable:
"""
Takes a nested iterable structure and applies f to all of its nodes,
returning them in the same structure.
"""

structure = type(nested_structure)
match nested_structure:
case dict():
return {
k: nested_map(f, v)
for k, v in nested_structure.items()
}
# NamedTuple
case _ if hasattr(structure, "_fields") and isinstance(
getattr(structure, "_fields"), tuple
):
d = {
k: nested_map(f, v) for
k, v in cast(NamedTuple, nested_structure)._asdict().items()
}
return cast(NestedIterable, structure(**d))
case _ if inspect.isgenerator(nested_structure):
return ( nested_map(f, s) for s in nested_structure )
# Have to check str first since it is an Iterable
case str():
return f(nested_structure)
case Iterable():
structure = cast(type

• , structure)
  return structure(map(lambda s: nested_map(f, s), nested_structure))
  case _:
  return f(nested_structure)
  
  
  Do you think there is a "more elegant" way to achieve this?
  
  Continue reading...