⚡️ Speed up function find_first_non_return_key by 42%

[codeflash-ai]

📄 42% (0.42x) speedup for find_first_non_return_key in gradio/cli/commands/components/_docs_utils.py

⏱️ Runtime : 64.0 microseconds → 45.0 microseconds (best of 75 runs)

📝 Explanation and details

The optimization replaces dictionary .items() iteration with direct key iteration, eliminating unnecessary tuple unpacking overhead.

Key changes:

• Changed for key, value in some_dict.items(): to for key in some_dict:
• Changed return value to return some_dict[key]

Why it's faster:

1. Eliminates tuple creation/unpacking: .items() creates (key, value) tuples for each iteration, which must then be unpacked. Direct key iteration avoids this allocation overhead.
2. Reduces function call overhead: .items() is a method call that returns an iterator, while some_dict iteration uses Python's built-in iterator protocol more efficiently.
3. Deferred value access: The optimized version only accesses some_dict[key] when actually returning, not for every iteration.

Performance impact:
The line profiler shows the loop line improved from 93,776ns to 79,529ns (15% faster), contributing to the overall 42% speedup. Test results show consistent 20-66% improvements across all scenarios, with particularly strong gains (50%+) for cases with empty dicts, single keys, or early returns.

Hot path impact:
Based on the function reference, this optimization is valuable since find_first_non_return_key is called within extract_docstrings() during component documentation processing. The function processes method parameters for each component's preprocess/postprocess methods, meaning this optimization reduces overhead during Gradio's CLI documentation generation workflow.

Best performance cases:
The optimization excels when the first non-"return" key is found early (up to 66% faster for empty dicts), making it particularly effective for typical documentation processing where "return" keys are less common or appear later in parameter dictionaries.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from gradio.cli.commands.components._docs_utils import \
    find_first_non_return_key

# unit tests

# 1. BASIC TEST CASES

def test_empty_dict_returns_none():
    # Test with an empty dictionary
    codeflash_output = find_first_non_return_key({}) # 921ns -> 553ns (66.5% faster)

def test_only_return_key():
    # Dict with only 'return' key should return None
    codeflash_output = find_first_non_return_key({"return": 123}) # 1.05μs -> 661ns (59.0% faster)

def test_return_and_one_other_key():
    # Dict with 'return' and one other key: should return value of the other key
    d = {"return": 1, "foo": 2}
    codeflash_output = find_first_non_return_key(d) # 1.07μs -> 794ns (35.4% faster)

def test_first_key_is_non_return():
    # Dict where first key is not 'return'
    d = {"bar": 7, "return": 8}
    codeflash_output = find_first_non_return_key(d) # 957ns -> 707ns (35.4% faster)

def test_multiple_non_return_keys():
    # Dict with several non-'return' keys
    d = {"a": 1, "b": 2, "return": 3, "c": 4}
    codeflash_output = find_first_non_return_key(d) # 1.00μs -> 665ns (51.0% faster)

def test_return_key_not_first():
    # 'return' key is not first; should return first non-'return' key's value
    d = {"x": 10, "return": 20, "y": 30}
    codeflash_output = find_first_non_return_key(d) # 1.01μs -> 695ns (44.9% faster)

def test_return_key_last():
    # 'return' key is last; should return first key's value
    d = {"foo": "bar", "baz": "qux", "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 981ns -> 659ns (48.9% faster)

def test_return_key_middle():
    # 'return' key in the middle
    d = {"foo": "bar", "return": "should not see", "baz": "qux"}
    codeflash_output = find_first_non_return_key(d) # 1.04μs -> 652ns (59.8% faster)

def test_multiple_return_keys():
    # Multiple 'return' keys, only first non-'return' key's value should be returned
    d = {"return": 1, "return": 2, "a": 3}
    # In Python, duplicate keys in dicts are overwritten; only last 'return' survives
    # So dict is {"return": 2, "a": 3}
    codeflash_output = find_first_non_return_key(d) # 1.09μs -> 749ns (45.1% faster)

# 2. EDGE TEST CASES

def test_key_is_return_with_different_case():
    # 'Return' key with different case should be treated as non-'return'
    d = {"Return": 5, "return": 6}
    codeflash_output = find_first_non_return_key(d) # 1.03μs -> 727ns (41.4% faster)

def test_key_is_return_with_spaces():
    # ' return ' key should be treated as non-'return'
    d = {" return ": 7, "return": 8}
    codeflash_output = find_first_non_return_key(d) # 977ns -> 713ns (37.0% faster)

def test_key_is_integer():
    # Non-string keys should not be equal to 'return'
    d = {1: "one", "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 1.06μs -> 821ns (29.5% faster)

def test_key_is_none():
    # None as key should be treated as non-'return'
    d = {None: "none", "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 1.10μs -> 845ns (30.1% faster)

def test_key_is_tuple():
    # Tuple as key
    d = {(1, 2): "tuple", "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 1.03μs -> 864ns (19.4% faster)

def test_key_is_empty_string():
    # Empty string key
    d = {"": "empty", "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 978ns -> 665ns (47.1% faster)

def test_key_is_false():
    # Boolean key
    d = {False: "false", "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 1.06μs -> 846ns (25.9% faster)

def test_key_is_return_unicode():
    # Unicode 'return' key
    d = {"rétúrn": "unicode", "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 1.03μs -> 762ns (35.6% faster)

def test_key_is_return_bytes():
    # Bytes key
    d = {b"return": "bytes", "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 1.01μs -> 776ns (30.0% faster)

def test_value_is_none():
    # Value is None for non-'return' key
    d = {"foo": None, "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 918ns -> 661ns (38.9% faster)

def test_value_is_false():
    # Value is False for non-'return' key
    d = {"foo": False, "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 942ns -> 645ns (46.0% faster)

def test_value_is_list():
    # Value is a list
    d = {"foo": [1,2,3], "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 984ns -> 675ns (45.8% faster)

def test_value_is_dict():
    # Value is a dict
    d = {"foo": {"bar": "baz"}, "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 957ns -> 683ns (40.1% faster)

def test_value_is_callable():
    # Value is a function
    def f(): return 42
    d = {"foo": f, "return": "should not see"}
    codeflash_output = find_first_non_return_key(d) # 944ns -> 682ns (38.4% faster)


def test_dict_with_only_return_key_and_none_value():
    # Only 'return' key, value is None
    d = {"return": None}
    codeflash_output = find_first_non_return_key(d) # 1.07μs -> 647ns (64.9% faster)

def test_dict_with_return_key_and_empty_string_value():
    # Only 'return' key, value is empty string
    d = {"return": ""}
    codeflash_output = find_first_non_return_key(d) # 1.03μs -> 646ns (58.8% faster)

def test_dict_with_return_key_and_false_value():
    # Only 'return' key, value is False
    d = {"return": False}
    codeflash_output = find_first_non_return_key(d) # 945ns -> 622ns (51.9% faster)

def test_dict_with_non_return_key_and_false_value():
    # Non-'return' key, value is False
    d = {"foo": False, "return": True}
    codeflash_output = find_first_non_return_key(d) # 972ns -> 708ns (37.3% faster)

# 3. LARGE SCALE TEST CASES

def test_large_dict_first_key_non_return():
    # Large dict, first key is non-'return'
    d = {"foo": 123}
    for i in range(1, 999):
        d[f"key{i}"] = i
    d["return"] = "should not see"
    codeflash_output = find_first_non_return_key(d) # 1.06μs -> 715ns (48.0% faster)

def test_large_dict_first_key_return():
    # Large dict, first key is 'return'
    d = {"return": 123}
    for i in range(1, 999):
        d[f"key{i}"] = i
    codeflash_output = find_first_non_return_key(d) # 1.10μs -> 749ns (47.0% faster)

def test_large_dict_all_return_keys():
    # Large dict, all keys are 'return'
    d = {"return": i for i in range(999)}
    codeflash_output = find_first_non_return_key(d) # 977ns -> 642ns (52.2% faster)

def test_large_dict_last_key_non_return():
    # Large dict, last key is non-'return'
    d = {"return": 0}
    for i in range(1, 998):
        d[f"return{i}"] = i  # These are not equal to "return"
    d["foo"] = "bar"
    codeflash_output = find_first_non_return_key(d) # 1.13μs -> 781ns (44.4% faster)

def test_large_dict_non_string_keys():
    # Large dict with integer keys
    d = {"return": "should not see"}
    for i in range(1, 999):
        d[i] = i
    codeflash_output = find_first_non_return_key(d) # 1.11μs -> 891ns (24.8% faster)

def test_large_dict_return_key_middle():
    # Large dict, 'return' key in the middle
    d = {}
    for i in range(500):
        d[f"key{i}"] = i
    d["return"] = "should not see"
    for i in range(500, 999):
        d[f"key{i}"] = i
    codeflash_output = find_first_non_return_key(d) # 973ns -> 666ns (46.1% faster)

def test_large_dict_return_key_last():
    # Large dict, 'return' key last
    d = {f"key{i}": i for i in range(999)}
    d["return"] = "should not see"
    codeflash_output = find_first_non_return_key(d) # 995ns -> 696ns (43.0% faster)

def test_large_dict_with_mixed_key_types():
    # Large dict with mixed key types
    d = {"return": "should not see"}
    for i in range(1, 500):
        d[i] = i
    for i in range(500, 999):
        d[str(i)] = i
    codeflash_output = find_first_non_return_key(d) # 1.22μs -> 938ns (29.5% faster)

def test_large_dict_with_none_key():
    # Large dict with None key
    d = {"return": "should not see"}
    for i in range(1, 998):
        d[f"key{i}"] = i
    d[None] = "none"
    codeflash_output = find_first_non_return_key(d) # 1.13μs -> 819ns (38.2% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
from gradio.cli.commands.components._docs_utils import \
    find_first_non_return_key

# unit tests

# --------------------------
# Basic Test Cases
# --------------------------

def test_single_non_return_key():
    # Only one key and it's not "return"
    d = {"foo": 42}
    codeflash_output = find_first_non_return_key(d) # 1.00μs -> 660ns (51.5% faster)

def test_return_key_first():
    # "return" is the first key, but there is another key
    d = {"return": "ignored", "bar": "found"}
    codeflash_output = find_first_non_return_key(d) # 1.08μs -> 725ns (49.7% faster)

def test_return_key_last():
    # "return" is the last key
    d = {"baz": 1, "return": 2}
    codeflash_output = find_first_non_return_key(d) # 999ns -> 695ns (43.7% faster)

def test_multiple_non_return_keys():
    # Multiple non-"return" keys, should return the first one
    d = {"a": 10, "b": 20, "return": 30, "c": 40}
    codeflash_output = find_first_non_return_key(d) # 993ns -> 679ns (46.2% faster)

def test_only_return_key():
    # Only "return" key present
    d = {"return": "only"}
    codeflash_output = find_first_non_return_key(d) # 983ns -> 621ns (58.3% faster)

# --------------------------
# Edge Test Cases
# --------------------------

def test_empty_dict():
    # Empty dictionary should return None
    d = {}
    codeflash_output = find_first_non_return_key(d) # 852ns -> 541ns (57.5% faster)

def test_return_key_with_none_value():
    # "return" key with None value, but another key exists
    d = {"return": None, "foo": "bar"}
    codeflash_output = find_first_non_return_key(d) # 1.06μs -> 797ns (32.9% faster)

def test_non_string_keys():
    # Keys are not strings, should still work
    d = {1: "one", "return": "ignored", 2: "two"}
    codeflash_output = find_first_non_return_key(d) # 1.07μs -> 834ns (28.2% faster)

def test_mixed_key_types():
    # Mixed key types, first non-"return" key is an int
    d = {"return": "skip", 0: "zero", "foo": "bar"}
    codeflash_output = find_first_non_return_key(d) # 1.16μs -> 900ns (28.9% faster)

def test_return_key_not_first():
    # "return" key is not first, should skip it
    d = {"x": 5, "return": 6, "y": 7}
    codeflash_output = find_first_non_return_key(d) # 1.01μs -> 686ns (46.8% faster)

def test_return_key_with_empty_string_value():
    # "return" key with empty string value, another key present
    d = {"return": "", "hello": "world"}
    codeflash_output = find_first_non_return_key(d) # 1.08μs -> 769ns (39.9% faster)

def test_all_keys_are_return():
    # All keys are "return"
    d = {"return": 1, "return": 2}
    # In Python, duplicate keys are not allowed, so only last survives
    codeflash_output = find_first_non_return_key(d) # 987ns -> 622ns (58.7% faster)

def test_key_named_return_but_not_string():
    # Key is not the string "return", but e.g. a tuple containing "return"
    d = {("return",): "tuple", "return": "string"}
    codeflash_output = find_first_non_return_key(d) # 1.08μs -> 855ns (26.9% faster)

def test_key_named_return_case_sensitive():
    # "Return" with different casing should not be skipped
    d = {"Return": "upper", "return": "lower"}
    codeflash_output = find_first_non_return_key(d) # 1.02μs -> 735ns (38.9% faster)

def test_key_named_return_with_spaces():
    # " return " with spaces should not be skipped
    d = {" return ": "spaced", "return": "exact"}
    codeflash_output = find_first_non_return_key(d) # 964ns -> 668ns (44.3% faster)

# --------------------------
# Large Scale Test Cases
# --------------------------

def test_large_dict_first_key_is_return():
    # Large dict with "return" as first key
    d = {"return": None}
    # Add 999 more keys
    for i in range(1, 1000):
        d[f"key{i}"] = i
    # Should return value of "key1"
    codeflash_output = find_first_non_return_key(d) # 1.12μs -> 791ns (41.5% faster)

def test_large_dict_return_key_in_middle():
    # Large dict with "return" in the middle
    d = {}
    for i in range(500):
        d[f"key{i}"] = i
    d["return"] = "skip"
    for i in range(500, 999):
        d[f"key{i}"] = i
    # Should return value of "key0"
    codeflash_output = find_first_non_return_key(d) # 978ns -> 722ns (35.5% faster)

def test_large_dict_all_return_keys():
    # Large dict with only "return" keys
    d = {"return": i for i in range(1000)}
    codeflash_output = find_first_non_return_key(d) # 967ns -> 657ns (47.2% faster)

def test_large_dict_non_return_key_last():
    # Large dict with "return" keys, last key is not "return"
    d = {"return": 0}
    for i in range(1, 999):
        d[f"return{i}"] = i  # These are not "return"
    d["not_return"] = "found"
    # Should return value of "return1" (first key not exactly "return")
    codeflash_output = find_first_non_return_key(d) # 1.17μs -> 765ns (52.8% faster)

def test_large_dict_first_key_not_return():
    # Large dict, first key is not "return"
    d = {"not_return": "first"}
    for i in range(1, 1000):
        d[f"key{i}"] = i
    d["return"] = "last"
    codeflash_output = find_first_non_return_key(d) # 1.01μs -> 679ns (49.2% faster)

def test_large_dict_with_non_string_keys():
    # Large dict with integer keys, and one "return" key
    d = {i: i for i in range(999)}
    d["return"] = "skip"
    codeflash_output = find_first_non_return_key(d) # 1.15μs -> 841ns (37.0% faster)

# --------------------------
# Mutation Testing Guards
# --------------------------

@pytest.mark.parametrize("d,expected", [
    # If function returns first key regardless of "return", it will fail these
    ({"return": "skip", "foo": "bar"}, "bar"),
    ({"foo": "bar", "return": "skip"}, "bar"),
    ({"return": "skip"}, None),
    ({}, None),
    ({"return": "skip", "return": "skip2"}, None),  # Only one key survives
])
def test_mutation_guard(d, expected):
    # Ensures that function only returns first non-"return" key
    codeflash_output = find_first_non_return_key(d) # 5.38μs -> 3.71μs (45.2% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-find_first_non_return_key-mhwv969h and push.