flatten_nested_dict

The oracle · bundled tests

"""
Tests for flatten_nested_dict module.
"""
import pytest
from flatten_nested_dict import flatten_nested_dict, unflatten_nested_dict


# ---------------------------------------------------------------------------
# flatten_nested_dict tests
# ---------------------------------------------------------------------------

class TestFlattenBasic:
    def test_empty_dict(self):
        assert flatten_nested_dict({}) == {}

    def test_single_level(self):
        assert flatten_nested_dict({"a": 1, "b": 2}) == {"a": 1, "b": 2}

    def test_two_levels(self):
        result = flatten_nested_dict({"a": {"b": 1}})
        assert result == {"a.b": 1}

    def test_three_levels(self):
        result = flatten_nested_dict({"a": {"b": {"c": 42}}})
        assert result == {"a.b.c": 42}

    def test_mixed_depth(self):
        data = {"x": 1, "y": {"z": 2, "w": {"v": 3}}}
        result = flatten_nested_dict(data)
        assert result == {"x": 1, "y.z": 2, "y.w.v": 3}

    def test_none_value_preserved(self):
        result = flatten_nested_dict({"a": {"b": None}})
        assert result == {"a.b": None}

    def test_false_value_preserved(self):
        result = flatten_nested_dict({"a": False})
        assert result == {"a": False}

    def test_zero_value_preserved(self):
        result = flatten_nested_dict({"a": {"b": 0}})
        assert result == {"a.b": 0}

    def test_empty_nested_dict_preserved(self):
        # An empty dict value should NOT be recursed into
        result = flatten_nested_dict({"a": {}})
        assert result == {"a": {}}


class TestFlattenSeparator:
    def test_custom_separator_slash(self):
        result = flatten_nested_dict({"a": {"b": 1}}, separator="/")
        assert result == {"a/b": 1}

    def test_custom_separator_double_underscore(self):
        result = flatten_nested_dict({"a": {"b": {"c": 99}}}, separator="__")

The artifact

"""
flatten_nested_dict.py

Flattens a nested dictionary into a single-level dictionary with
dot-separated keys (or a custom separator). Supports an optional
maximum depth limit and handles lists by indexing them.
"""


def flatten_nested_dict(
    data: dict,
    separator: str = ".",
    max_depth: int = None,
    _prefix: str = "",
    _current_depth: int = 0,
) -> dict:
    """
    Flatten a nested dictionary into a single-level dict.

    Args:
        data: The dictionary to flatten.
        separator: The string used to join nested keys. Defaults to ".".
        max_depth: Maximum depth to flatten. None means unlimited.
                   At max_depth, values are left as-is (not further flattened).
        _prefix: Internal use — the key prefix accumulated so far.
        _current_depth: Internal use — tracks current recursion depth.

    Returns:
        A flat dictionary with compound keys.

    Examples:
        >>> flatten_nested_dict({"a": {"b": {"c": 1}}})
        {'a.b.c': 1}

        >>> flatten_nested_dict({"a": {"b": 1}, "c": 2}, separator="/")
        {'a/b': 1, 'c': 2}

        >>> flatten_nested_dict({"a": {"b": {"c": 1}}}, max_depth=1)
        {'a.b': {'c': 1}}
    """
    if not isinstance(data, dict):
        raise TypeError(f"Expected a dict, got {type(data).__name__!r}")

    result = {}

    for key, value in data.items():
        str_key = str(key)
        full_key = f"{_prefix}{separator}{str_key}" if _prefix else str_key

        depth_exceeded = max_depth is not None and _current_depth >= max_depth

        if isinstance(value, dict) and value and not depth_exceeded:
            nested = flatten_nested_dict(
                value,
                separator=separator,
                max_depth=max_depth,
                _prefix=full_key,
                _current_depth=_current_depth + 1,
            )
            result.update(nested)
        elif isinstance(value, list) and not depth_exceeded:
            for i, item in enumerate(value):
                item_key = f"{full_key}{separator}{i}"
                if isinstance(item, dict) and item:
                    nested = flatten_nested_dict(
                        item,
                        separator=separator,
                        max_depth=max_depth,
                        _prefix=item_key,
                        _current_depth=_current_depth + 1,
                    )
                    result.update(nested)
                else:
                    result[item_key] = item
        else:
            result[full_key] = value

    return result


def unflatten_nested_dict(data: dict, separator: str = ".") -> dict:
    """
    Reverse a flattened dictionary back into a nested dictionary.

    Note: List reconstruction is not supported; numeric segments are
    treated as regular string keys.

    Args:
        data: A flat dictionary with compound keys.
        separator: The separator used in the compound keys. Defaults to ".".

    Returns:
        A nested dictionary.

    Examples:
        >>> unflatten_nested_dict({"a.b.c": 1, "a.b.d": 2})
        {'a': {'b': {'c': 1, 'd': 2}}}
    """
    result = {}

    for compound_key, value in data.items():
        parts = compound_key.split(separator)
        node = result
        for part in parts[:-1]:
            if part not in node or not isinstance(node[part], dict):
                node[part] = {}
            node = node[part]
        node[parts[-1]] = value

    return result