⚡️ Speed up function flatten_dict by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for flatten_dict in unstructured/staging/base.py

⏱️ Runtime : 6.95 milliseconds → 6.16 milliseconds (best of 30 runs)

📝 Explanation and details

The optimized code achieves a 12% speedup by restructuring the recursion strategy to eliminate repeated dictionary creation and .update() calls.

Key Optimization

Original approach: Creates temporary dictionaries and merges them via .update() for each recursive call:

flattened_dict.update(
    flatten_dict(value, new_key, separator, ...)
)

Optimized approach: Uses a nested helper function _flatten_recursive() that writes directly to the shared flattened_dict:

def _flatten_recursive(value: Any, current_key: str) -> None:
    # Directly writes to flattened_dict
    flattened_dict[current_key] = value

Why This is Faster

1. Eliminates dictionary allocation overhead: The original code creates a new dictionary for every recursive call (4,332 calls per profiler), then merges it. The optimized version writes directly to one dictionary.

2. Avoids .update() operations: Dictionary merging is expensive - it involves iterating through key-value pairs and copying them. The line profiler shows this accounts for ~10% of total time in the original code.

3. Reduces function call overhead: While both versions recurse, the optimized version doesn't pass dictionaries as arguments or return values in recursive calls, reducing parameter passing overhead.

4. Simpler list/tuple handling: The original wraps list items in temporary dictionaries ({f"{new_key}{separator}{index}": item}), which the optimized version avoids entirely.

Performance Impact by Workload

Based on annotated tests:

• Large lists/nested structures (hot path): 30-68% faster - significant wins when flatten_lists=True on large datasets (test cases like test_flatten_large_list_flattening, test_flatten_mixed_large_structure)
• Deep nesting: 25% faster - benefits from eliminating repeated dictionary creation in deep recursion
• Simple/flat dictionaries: 10-30% slower - slight overhead from the helper function closure when recursion depth is minimal

Function References Context

The function is called in convert_to_csv() and convert_to_dataframe() to flatten metadata dictionaries. These appear to be data transformation pipelines where:

• Elements are processed in batches (loops over rows in CSV conversion)
• Each call flattens potentially nested metadata structures

The optimization is particularly beneficial here because metadata flattening happens per-row in batch operations, making the cumulative effect of the 12% speedup meaningful in data-heavy workloads.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 37 Passed
🌀 Generated Regression Tests	✅ 62 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	✅ 3 Passed
📊 Tests Coverage	100.0%

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
staging/test_base.py::test_default_pandas_dtypes	20.2μs	23.3μs	-13.5%⚠️
staging/test_base.py::test_flatten_dict	3.61μs	5.20μs	-30.6%⚠️
staging/test_base.py::test_flatten_dict_alt_separator	6.72μs	7.66μs	-12.3%⚠️
staging/test_base.py::test_flatten_dict_empty_lists	4.99μs	5.63μs	-11.3%⚠️
staging/test_base.py::test_flatten_dict_flatten_list	14.3μs	12.6μs	13.4%✅
staging/test_base.py::test_flatten_dict_flatten_list_none_in_list_remove_none	13.5μs	12.1μs	11.5%✅
staging/test_base.py::test_flatten_dict_flatten_list_omit_keys	9.84μs	9.34μs	5.33%✅
staging/test_base.py::test_flatten_dict_flatten_list_omit_keys2	10.6μs	10.5μs	1.02%✅
staging/test_base.py::test_flatten_dict_flatten_list_omit_keys3	11.7μs	11.2μs	5.11%✅
staging/test_base.py::test_flatten_dict_flatten_list_omit_keys4	12.0μs	11.0μs	8.91%✅
staging/test_base.py::test_flatten_dict_flatten_list_omit_keys_remove_none	9.60μs	9.16μs	4.87%✅
staging/test_base.py::test_flatten_dict_flatten_list_remove_none	13.9μs	12.0μs	15.3%✅
staging/test_base.py::test_flatten_dict_flatten_tuple	14.2μs	12.6μs	12.2%✅
staging/test_base.py::test_flatten_dict_with_lists	6.57μs	7.06μs	-6.98%⚠️
staging/test_base.py::test_flatten_dict_with_omit_keys	4.00μs	5.77μs	-30.7%⚠️
staging/test_base.py::test_flatten_dict_with_tuples	6.77μs	7.41μs	-8.69%⚠️
staging/test_base.py::test_flatten_empty_dict	1.44μs	2.11μs	-31.9%⚠️
staging/test_base.py::test_flatten_nested_dict	6.44μs	7.22μs	-10.7%⚠️

🌀 Click to see Generated Regression Tests

from __future__ import annotations

# imports
from unstructured.staging.base import flatten_dict


def test_basic_flattening_simple_nested_dict():
    # Basic case: a simple nested dict should be flattened with default separator '_'
    data = {"a": {"b": 1}, "c": 2}
    codeflash_output = flatten_dict(data)
    result = codeflash_output  # 5.42μs -> 6.12μs (11.4% slower)


def test_custom_separator_and_parent_key_behavior():
    # Using a custom separator and providing a parent_key should produce concatenated keys accordingly
    data = {"b": 1}
    # parent_key 'a' and separator '.' should make key 'a.b'
    codeflash_output = flatten_dict(data, parent_key="a", separator=".")
    result = codeflash_output  # 3.47μs -> 4.96μs (30.0% slower)


def test_flatten_lists_true_flattens_list_of_dicts_and_primitives():
    # When flatten_lists is True, lists of dicts become flattened with indexed keys
    data = {"lst": [{"x": 1}, {"y": 2}], "prims": [10, 20]}
    codeflash_output = flatten_dict(data, flatten_lists=True)
    result = codeflash_output  # 12.9μs -> 11.8μs (8.70% faster)
    # Expect 'lst_0_x', 'lst_1_y' for the dict items and 'prims_0', 'prims_1' for primitives
    expected = {"lst_0_x": 1, "lst_1_y": 2, "prims_0": 10, "prims_1": 20}


def test_flatten_lists_false_preserves_lists_intact():
    # When flatten_lists is False, lists should not be broken into indexed keys
    data = {"lst": [{"x": 1}, {"y": 2}], "prims": [10, 20]}
    codeflash_output = flatten_dict(data, flatten_lists=False)
    result = codeflash_output  # 3.17μs -> 4.83μs (34.3% slower)


def test_remove_none_filters_none_values_recursively():
    # remove_none True should remove entries whose value is None at any level
    data = {"a": None, "b": {"c": None, "d": 3}}
    codeflash_output = flatten_dict(data, remove_none=True)
    result = codeflash_output  # 5.55μs -> 6.75μs (17.8% slower)

    # With remove_none False, None values should remain in the flattened output
    codeflash_output = flatten_dict(data, remove_none=False)
    result_no_remove = codeflash_output  # 3.73μs -> 4.10μs (8.93% slower)


def test_keys_to_omit_prevents_flattening_of_named_subtrees():
    # keys_to_omit contains the full new_key (parent + separator + key) that should not be flattened
    data = {"parent": {"omit": {"inner": 1}, "keep": 2}}
    # Omit the 'parent_omit' path so the nested dict at that path is preserved as-is
    codeflash_output = flatten_dict(data, keys_to_omit=["parent_omit"])
    result = codeflash_output  # 5.64μs -> 6.59μs (14.4% slower)


def test_keys_to_omit_with_lists_and_flattening():
    # When flattening lists, the produced keys include the index suffix; keys_to_omit can match those keys
    data = {"arr": [{"a": 1}, {"b": 2}]}
    # If we omit 'arr_0', the first item should be kept as a dict under that key and second flattened
    codeflash_output = flatten_dict(data, flatten_lists=True, keys_to_omit=["arr_0"])
    result = codeflash_output  # 9.33μs -> 8.74μs (6.74% faster)


def test_tuple_flattening_behaves_like_list_when_enabled():
    # Tuples should be treated similar to lists when flatten_lists=True
    data = {"t": ({"a": 1}, {"b": 2})}
    codeflash_output = flatten_dict(data, flatten_lists=True)
    result = codeflash_output  # 10.2μs -> 9.43μs (8.34% faster)


def test_empty_dict_and_empty_string_key_behavior():
    # Empty dictionary should return an empty dict
    codeflash_output = flatten_dict({})  # 1.51μs -> 2.21μs (31.9% slower)

    # If a top-level key is an empty string and its value is a dict, the empty parent_key
    # is falsy, so nested keys should not be prefixed by an extra separator.
    data = {"": {"a": 1}}
    codeflash_output = flatten_dict(data)
    result = codeflash_output  # 3.59μs -> 3.80μs (5.57% slower)


def test_non_string_top_level_key_casting_and_resulting_key_names():
    # While typing suggests str keys, the function uses f-strings and will coerce non-str parent_key to string
    data = {1: {"a": 2}}  # integer top-level key
    codeflash_output = flatten_dict(data)
    result = codeflash_output  # 4.89μs -> 5.56μs (12.0% slower)


def test_large_scale_flattening_many_items():
    # Large-scale test: flatten a wide structure with many entries to test performance and correctness
    # Use N significantly less than 1000 to comply with requirements and keep tests fast/deterministic
    N = 500
    # Build a dictionary with N top-level keys each containing a small nested dict
    big = {f"k{i}": {"sub": i} for i in range(N)}
    codeflash_output = flatten_dict(big)
    result = codeflash_output  # 486μs -> 452μs (7.50% faster)


def test_mixed_depth_and_types_combination():
    # Complex case mixing dicts, lists, None, and keys_to_omit to exercise multiple branches
    data = {
        "a": {"b": {"c": None, "d": 4}, "e": [1, {"f": 6}]},
        "g": None,
        "h": (7, {"i": 8}),
    }
    # Flatten lists/tuples and remove None values, but omit 'a_b' subtree so it remains intact
    codeflash_output = flatten_dict(
        data, flatten_lists=True, remove_none=True, keys_to_omit=["a_b"]
    )
    result = codeflash_output  # 15.7μs -> 14.3μs (9.99% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from unstructured.staging.base import flatten_dict


def test_flatten_simple_nested_dict():
    """Test flattening a simple two-level nested dictionary."""
    input_dict = {"a": {"b": 1, "c": 2}}
    expected = {"a_b": 1, "a_c": 2}
    codeflash_output = flatten_dict(input_dict)  # 5.70μs -> 6.41μs (11.1% slower)


def test_flatten_flat_dict():
    """Test that a flat dictionary remains unchanged."""
    input_dict = {"a": 1, "b": 2, "c": 3}
    expected = {"a": 1, "b": 2, "c": 3}
    codeflash_output = flatten_dict(input_dict)  # 3.81μs -> 5.24μs (27.2% slower)


def test_flatten_deeply_nested_dict():
    """Test flattening a deeply nested dictionary with multiple levels."""
    input_dict = {"a": {"b": {"c": {"d": 1}}}}
    expected = {"a_b_c_d": 1}
    codeflash_output = flatten_dict(input_dict)  # 6.97μs -> 6.86μs (1.56% faster)


def test_flatten_with_custom_separator():
    """Test flattening with a custom separator character."""
    input_dict = {"a": {"b": 1, "c": 2}}
    expected = {"a.b": 1, "a.c": 2}
    codeflash_output = flatten_dict(input_dict, separator=".")  # 5.77μs -> 6.71μs (14.0% slower)


def test_flatten_mixed_types():
    """Test flattening a dictionary with mixed value types."""
    input_dict = {"a": {"b": 1, "c": "string", "d": 3.14, "e": True}}
    expected = {"a_b": 1, "a_c": "string", "a_d": 3.14, "a_e": True}
    codeflash_output = flatten_dict(input_dict)  # 6.68μs -> 7.50μs (10.9% slower)


def test_flatten_empty_dict():
    """Test flattening an empty dictionary."""
    input_dict = {}
    expected = {}
    codeflash_output = flatten_dict(input_dict)  # 1.51μs -> 2.21μs (31.6% slower)


def test_flatten_empty_nested_dict():
    """Test flattening a dictionary with empty nested dictionaries."""
    input_dict = {"a": {}, "b": 1}
    expected = {"b": 1}
    codeflash_output = flatten_dict(input_dict)  # 4.62μs -> 4.98μs (7.07% slower)


def test_flatten_with_none_values_default():
    """Test that None values are kept by default."""
    input_dict = {"a": {"b": None, "c": 1}}
    expected = {"a_b": None, "a_c": 1}
    codeflash_output = flatten_dict(input_dict)  # 5.63μs -> 6.29μs (10.6% slower)


def test_flatten_with_none_values_removed():
    """Test that None values are removed when remove_none is True."""
    input_dict = {"a": {"b": None, "c": 1}}
    expected = {"a_c": 1}
    codeflash_output = flatten_dict(input_dict, remove_none=True)  # 5.70μs -> 6.36μs (10.4% slower)


def test_flatten_dict_with_string_keys():
    """Test flattening with various string key formats."""
    input_dict = {"key1": {"nested_key": 1, "another_key": 2}}
    expected = {"key1_nested_key": 1, "key1_another_key": 2}
    codeflash_output = flatten_dict(input_dict)  # 5.68μs -> 6.32μs (10.1% slower)


def test_flatten_with_numeric_string_keys():
    """Test flattening with numeric strings as keys."""
    input_dict = {"1": {"2": {"3": "value"}}}
    expected = {"1_2_3": "value"}
    codeflash_output = flatten_dict(input_dict)  # 5.75μs -> 5.88μs (2.09% slower)


def test_flatten_omit_single_key():
    """Test omitting a single key from flattening."""
    input_dict = {"a": {"b": 1, "c": 2}}
    expected = {"a_b": 1, "a_c": 2}
    codeflash_output = flatten_dict(input_dict, keys_to_omit=["a_b"])
    result = codeflash_output  # 5.86μs -> 6.65μs (11.9% slower)


def test_flatten_omit_key_prevents_flattening():
    """Test that omitted keys are not flattened further."""
    input_dict = {"a": {"b": {"c": 1}}}
    expected = {"a_b": {"c": 1}}
    codeflash_output = flatten_dict(input_dict, keys_to_omit=["a_b"])
    result = codeflash_output  # 4.37μs -> 5.13μs (14.9% slower)


def test_flatten_omit_multiple_keys():
    """Test omitting multiple keys from flattening."""
    input_dict = {"a": {"b": 1, "c": 2, "d": 3}}
    expected = {"a_b": 1, "a_c": 2, "a_d": 3}
    codeflash_output = flatten_dict(input_dict, keys_to_omit=["a_b", "a_c"])
    result = codeflash_output  # 6.07μs -> 7.19μs (15.7% slower)


def test_flatten_with_list_flatten_disabled():
    """Test that lists are not flattened when flatten_lists is False."""
    input_dict = {"a": {"b": [1, 2, 3]}}
    expected = {"a_b": [1, 2, 3]}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=False
    )  # 4.89μs -> 5.71μs (14.3% slower)


def test_flatten_with_list_flatten_enabled():
    """Test that lists are flattened when flatten_lists is True."""
    input_dict = {"a": {"b": [1, 2, 3]}}
    expected = {"a_b_0": 1, "a_b_1": 2, "a_b_2": 3}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=True
    )  # 10.2μs -> 9.40μs (8.38% faster)


def test_flatten_with_tuple_flatten_enabled():
    """Test that tuples are flattened when flatten_lists is True."""
    input_dict = {"a": {"b": (1, 2, 3)}}
    expected = {"a_b_0": 1, "a_b_1": 2, "a_b_2": 3}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=True
    )  # 10.2μs -> 9.30μs (9.81% faster)


def test_flatten_with_empty_list():
    """Test flattening with an empty list."""
    input_dict = {"a": {"b": []}}
    expected = {"a_b": []}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=False
    )  # 4.84μs -> 5.77μs (16.2% slower)


def test_flatten_with_nested_list_of_dicts():
    """Test flattening lists containing dictionaries when flatten_lists is True."""
    input_dict = {"a": [{"b": 1}, {"c": 2}]}
    expected = {"a_0_b": 1, "a_1_c": 2}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=True
    )  # 10.4μs -> 9.61μs (8.45% faster)


def test_flatten_with_nested_dict_in_list():
    """Test that nested dicts within lists are properly flattened."""
    input_dict = {"items": [{"nested": {"value": 1}}]}
    expected = {"items_0_nested_value": 1}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=True
    )  # 9.05μs -> 8.49μs (6.55% faster)


def test_flatten_with_none_in_list():
    """Test that None values in lists are handled correctly."""
    input_dict = {"a": [1, None, 3]}
    expected = {"a_0": 1, "a_1": None, "a_2": 3}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=True
    )  # 9.10μs -> 8.37μs (8.66% faster)


def test_flatten_with_none_in_list_and_remove_none():
    """Test that None values in lists are removed when remove_none is True."""
    input_dict = {"a": [1, None, 3]}
    expected = {"a_0": 1, "a_2": 3}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=True, remove_none=True
    )  # 8.99μs -> 8.14μs (10.4% faster)


def test_flatten_special_characters_in_keys():
    """Test flattening with special characters in dictionary keys."""
    input_dict = {"key-1": {"nested_key": 1}}
    expected = {"key-1_nested_key": 1}
    codeflash_output = flatten_dict(input_dict)  # 4.83μs -> 5.25μs (7.91% slower)


def test_flatten_unicode_keys():
    """Test flattening with unicode characters in keys."""
    input_dict = {"keyo": {"key": 1}}
    expected = {"keyo_key": 1}
    codeflash_output = flatten_dict(input_dict)  # 4.86μs -> 5.40μs (10.1% slower)


def test_flatten_all_none_values_with_remove_none():
    """Test that dictionary becomes empty when all values are None and remove_none is True."""
    input_dict = {"a": {"b": None, "c": None}}
    expected = {}
    codeflash_output = flatten_dict(input_dict, remove_none=True)  # 4.55μs -> 5.37μs (15.2% slower)


def test_flatten_mix_remove_none_and_keep_omitted():
    """Test that omitted keys are kept even if they are None when remove_none is True."""
    input_dict = {"a": {"b": None, "c": 1}}
    expected = {"a_b": None, "a_c": 1}
    codeflash_output = flatten_dict(input_dict, remove_none=True, keys_to_omit=["a_b"])
    result = codeflash_output  # 5.65μs -> 6.70μs (15.6% slower)


def test_flatten_with_parent_key_parameter():
    """Test that parent_key parameter is used correctly."""
    input_dict = {"b": 1, "c": 2}
    expected = {"a_b": 1, "a_c": 2}
    codeflash_output = flatten_dict(input_dict, parent_key="a")  # 4.25μs -> 5.69μs (25.4% slower)


def test_flatten_complex_separator():
    """Test flattening with a multi-character separator."""
    input_dict = {"a": {"b": {"c": 1}}}
    expected = {"a::b::c": 1}
    codeflash_output = flatten_dict(input_dict, separator="::")  # 6.18μs -> 6.65μs (7.00% slower)


def test_flatten_with_boolean_values():
    """Test flattening with boolean values."""
    input_dict = {"a": {"b": True, "c": False}}
    expected = {"a_b": True, "a_c": False}
    codeflash_output = flatten_dict(input_dict)  # 5.40μs -> 6.02μs (10.3% slower)


def test_flatten_with_zero_values():
    """Test that zero values are preserved and not treated as None."""
    input_dict = {"a": {"b": 0, "c": 0.0}}
    expected = {"a_b": 0, "a_c": 0.0}
    codeflash_output = flatten_dict(input_dict, remove_none=True)  # 5.88μs -> 6.79μs (13.5% slower)


def test_flatten_with_empty_string_values():
    """Test that empty string values are preserved."""
    input_dict = {"a": {"b": "", "c": "value"}}
    expected = {"a_b": "", "a_c": "value"}
    codeflash_output = flatten_dict(input_dict)  # 5.22μs -> 5.91μs (11.8% slower)


def test_flatten_with_list_of_strings():
    """Test flattening a list of strings."""
    input_dict = {"a": ["x", "y", "z"]}
    expected = {"a_0": "x", "a_1": "y", "a_2": "z"}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=True
    )  # 8.80μs -> 8.09μs (8.71% faster)


def test_flatten_nested_lists():
    """Test flattening nested lists."""
    input_dict = {"a": [[1, 2], [3, 4]]}
    expected = {"a_0_0": 1, "a_0_1": 2, "a_1_0": 3, "a_1_1": 4}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=True
    )  # 13.1μs -> 10.9μs (19.3% faster)


def test_flatten_with_omit_and_flatten_lists():
    """Test omitting keys while also flattening lists."""
    input_dict = {"a": [{"b": 1}, {"c": 2}]}
    expected = {"a_0_b": 1, "a_1_c": 2}
    codeflash_output = flatten_dict(input_dict, flatten_lists=True, keys_to_omit=["nonexistent"])
    result = codeflash_output  # 10.3μs -> 9.74μs (6.19% faster)


def test_flatten_preserves_non_dict_types_in_values():
    """Test that non-dict types in values are preserved exactly."""
    input_dict = {"a": {"b": (1, 2, 3)}}
    expected = {"a_b": (1, 2, 3)}
    codeflash_output = flatten_dict(
        input_dict, flatten_lists=False
    )  # 5.19μs -> 5.76μs (9.88% slower)


def test_flatten_with_list_containing_none_elements():
    """Test flattening with None elements in lists."""
    input_dict = {"a": [1, None, 2]}
    expected = {"a_0": 1, "a_1": None, "a_2": 2}
    codeflash_output = flatten_dict(input_dict, flatten_lists=True, remove_none=False)
    result = codeflash_output  # 9.46μs -> 8.43μs (12.2% faster)


def test_flatten_large_flat_dictionary():
    """Test flattening a large flat dictionary."""
    input_dict = {f"key_{i}": i for i in range(500)}
    codeflash_output = flatten_dict(input_dict)
    result = codeflash_output  # 149μs -> 196μs (24.2% slower)


def test_flatten_large_nested_dictionary():
    """Test flattening a large nested dictionary."""
    # Create a dictionary with 100 top-level keys, each containing 5 nested keys
    input_dict = {f"parent_{i}": {f"child_{j}": i * 10 + j for j in range(5)} for i in range(100)}
    codeflash_output = flatten_dict(input_dict)
    result = codeflash_output  # 280μs -> 298μs (6.18% slower)


def test_flatten_deeply_nested_large_dictionary():
    """Test flattening a deeply nested dictionary structure."""
    # Create a dictionary with 5 levels of nesting
    input_dict = {"level_0": {"level_1": {"level_2": {"level_3": {"level_4": {"value": 42}}}}}}
    codeflash_output = flatten_dict(input_dict)
    result = codeflash_output  # 8.68μs -> 8.39μs (3.38% faster)


def test_flatten_large_list_of_dicts():
    """Test flattening a large list of dictionaries."""
    input_dict = {"items": [{"id": i, "value": i * 2} for i in range(200)]}
    codeflash_output = flatten_dict(input_dict, flatten_lists=True)
    result = codeflash_output  # 418μs -> 321μs (30.3% faster)


def test_flatten_large_dict_with_remove_none():
    """Test flattening a large dictionary with many None values."""
    # Create a dictionary where approximately half the values are None
    input_dict = {
        f"parent_{i}": {f"child_{j}": j if j % 2 == 0 else None for j in range(50)}
        for i in range(50)
    }
    codeflash_output = flatten_dict(input_dict, remove_none=True)
    result = codeflash_output  # 804μs -> 1.01ms (20.0% slower)


def test_flatten_large_dict_with_custom_separator():
    """Test flattening a large nested dictionary with custom separator."""
    input_dict = {f"key_{i}": {f"nested_{j}": i * j for j in range(50)} for i in range(50)}
    codeflash_output = flatten_dict(input_dict, separator="::")
    result = codeflash_output  # 1.18ms -> 1.30ms (8.68% slower)


def test_flatten_large_list_flattening():
    """Test flattening a large dictionary with list values."""
    input_dict = {"data": [i for i in range(300)]}
    codeflash_output = flatten_dict(input_dict, flatten_lists=True)
    result = codeflash_output  # 274μs -> 171μs (60.1% faster)


def test_flatten_mixed_large_structure():
    """Test flattening a large mixed structure with dicts and lists."""
    input_dict = {
        f"group_{i}": {f"subgroup_{j}": [k for k in range(5)] for j in range(10)} for i in range(10)
    }
    codeflash_output = flatten_dict(input_dict, flatten_lists=True)
    result = codeflash_output  # 542μs -> 349μs (55.1% faster)


def test_flatten_large_dict_with_omit_keys():
    """Test flattening a large dictionary while omitting specific keys."""
    input_dict = {f"key_{i}": {f"nested_{j}": i * j for j in range(20)} for i in range(20)}
    omit_keys = [f"key_{i}_nested_0" for i in range(20)]  # Omit 20 keys
    codeflash_output = flatten_dict(input_dict, keys_to_omit=omit_keys)
    result = codeflash_output  # 310μs -> 328μs (5.43% slower)


def test_flatten_large_dict_all_parameters():
    """Test flattening with all parameters on a large dictionary."""
    input_dict = {
        f"group_{i}": {
            f"item_{j}": [k if k % 3 != 0 else None for k in range(10)] for j in range(15)
        }
        for i in range(10)
    }
    codeflash_output = flatten_dict(
        input_dict,
        flatten_lists=True,
        remove_none=True,
        separator=".",
    )
    result = codeflash_output  # 1.36ms -> 808μs (68.8% faster)


def test_flatten_extreme_depth():
    """Test flattening with very deep nesting."""
    # Build a deeply nested dictionary
    input_dict = {"a": 1}
    for i in range(50):
        input_dict = {f"level_{i}": input_dict}
    codeflash_output = flatten_dict(input_dict)
    result = codeflash_output  # 51.1μs -> 40.6μs (25.6% faster)


def test_flatten_large_list_with_mixed_types():
    """Test flattening a large list containing mixed types."""
    input_dict = {
        "mixed_list": [1, "string", 3.14, True, None, {"nested": "dict"}, [1, 2, 3], "another"]
        * 50  # Repeat 50 times
    }
    codeflash_output = flatten_dict(input_dict, flatten_lists=True, remove_none=True)
    result = codeflash_output  # 551μs -> 343μs (60.5% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from pandas._libs.tslibs.strptime import TimeRE

from unstructured.staging.base import flatten_dict


def test_flatten_dict():
    flatten_dict(
        {"": [{}]},
        parent_key="",
        separator="",
        flatten_lists=True,
        remove_none=False,
        keys_to_omit=None,
    )


def test_flatten_dict_2():
    flatten_dict(
        {"": [], "\x00": TimeRE(locale_time="")},
        parent_key="",
        separator="",
        flatten_lists=False,
        remove_none=False,
        keys_to_omit=None,
    )


def test_flatten_dict_3():
    flatten_dict(
        {"": {0: {"is_ordered": False, "is_dictionary": False, "categories": None}}},
        parent_key="",
        separator="",
        flatten_lists=False,
        remove_none=False,
        keys_to_omit=None,
    )

🔎 Click to see Concolic Coverage Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_xdo_puqm/tmphc_5lxaq/test_concolic_coverage.py::test_flatten_dict	7.01μs	6.66μs	5.19%✅
codeflash_concolic_xdo_puqm/tmphc_5lxaq/test_concolic_coverage.py::test_flatten_dict_2	18.3μs	20.7μs	-11.5%⚠️
codeflash_concolic_xdo_puqm/tmphc_5lxaq/test_concolic_coverage.py::test_flatten_dict_3	6.50μs	7.21μs	-9.83%⚠️

To edit these changes git checkout codeflash/optimize-flatten_dict-mks0isd1 and push.