flatten_nested.py

from typing import Dict, Iterator, Tuple, TypeVar, Union

K = TypeVar("K")
V = TypeVar("V")
NestedDict = Union[V, Dict[K, Union["NestededDict", V]]]


def flatten_nested(
    nested: NestedDict, parent_keys: Tuple[K, ...] = ()
) -> Iterator[Tuple[Tuple[K, ...], V]]:
    if isinstance(nested, dict):
        for key, value in nested.items():
            keys = parent_keys + (key,)
            yield from flatten_nested(value, parent_keys=keys)
    else:
        yield parent_keys, nested


if __name__ == "__main__":
    for nested, expected in [
        ("foo", [((), "foo")]),
        ({"a": "foo"}, [(("a",), "foo")]),
        ({"a": {"b": "foo"}}, [(("a", "b"), "foo")]),
    ]:
        assert list(flatten_nested(nested)) == expected