⚡️ Speed up function extract_docstrings by 8,281%

[codeflash-ai]

📄 8,281% (82.81x) speedup for extract_docstrings in gradio/cli/commands/components/_docs_utils.py

⏱️ Runtime : 1.56 seconds → 18.7 milliseconds (best of 74 runs)

📝 Explanation and details

The optimized code achieves an 8281% speedup by targeting the primary performance bottleneck in the add_value function, which was consuming over 99% of execution time by repeatedly calling the expensive format() function.

Key optimizations:

1. Memoization in add_value: Added a function-level cache (_format_cache) that stores results of the slow format() calls. Since format() is deterministic for identical inputs, this dramatically reduces redundant subprocess invocations to the ruff formatter. The cache achieves a ~77% hit rate (367 cache hits out of 474 calls), eliminating most expensive operations.

2. Micro-optimizations:

   • find_first_non_return_key: Iterates only over keys instead of key-value pairs, accessing values only when needed
   • set_deep: Uses fast .get() check before .setdefault() to avoid unnecessary dict creation
   • extract_docstrings: Pre-compiles regex and filters lines containing ":" before applying expensive regex matching

Performance impact: The test results show consistent 4000-8000% improvements across various workloads, with the largest gains in scenarios with many classes and methods where add_value is called repeatedly. Based on the function reference, this optimization significantly benefits the CLI documentation generation workflow in gradio/cli/commands/components/docs.py, where extract_docstrings is called during the main documentation pipeline. The speedup transforms what was likely a multi-second operation into a sub-second one, dramatically improving developer experience when generating component documentation.

The optimization preserves all behavioral contracts - the memoization only caches deterministic results, and all micro-optimizations maintain identical logic paths and return values.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import inspect
import re
import sys
import textwrap
import types
import typing
from subprocess import PIPE, Popen

# imports
import pytest
from gradio.cli.commands.components._docs_utils import extract_docstrings

# -------------------
# UNIT TESTS BELOW
# -------------------

# Utilities for tests
def make_module(name, members):
    """Dynamically create a module with given members."""
    mod = types.ModuleType(name)
    for k, v in members.items():
        setattr(mod, k, v)
    return mod

# -------------------
# 1. BASIC TEST CASES
# -------------------




def test_extract_docstrings_multiple_classes():
    """Test extraction from module with multiple classes."""

    class C1:
        """C1 doc"""
        def __init__(self, a: int): """a: doc"""; pass
        def preprocess(self, a: int) -> int: """a: doc\nReturns: doc"""; return a

    class C2:
        """C2 doc"""
        def __init__(self, b: int): """b: doc"""; pass
        def preprocess(self, b: int) -> int: """b: doc\nReturns: doc"""; return b

    mod = make_module("mod4", {"C1": C1, "C2": C2})
    docs, mode = extract_docstrings(mod) # 13.6ms -> 269μs (4952% faster)

# -------------------
# 2. EDGE TEST CASES
# -------------------

def test_extract_docstrings_empty_module():
    """Test extraction from an empty module."""
    mod = make_module("mod_empty", {})
    docs, mode = extract_docstrings(mod) # 8.71μs -> 9.14μs (4.75% slower)

def test_extract_docstrings_class_no_methods():
    """Test extraction from a class with no __init__, preprocess, or postprocess."""
    class NoMethods:
        """No methods class doc"""
        pass
    mod = make_module("mod_nomethods", {"NoMethods": NoMethods})
    docs, mode = extract_docstrings(mod) # 44.1μs -> 40.5μs (8.89% faster)



def test_extract_docstrings_class_with_events():
    """Test extraction from a class with EVENTS list."""

    class Event:
        def __init__(self, name, doc):
            self.name = name
            self.doc = doc
        def __str__(self): return self.name

    class Dummy6:
        """Has events"""
        EVENTS = [Event("foo", "foo doc for {{ component }}"), Event("bar", "bar doc for {{ component }}")]
        def __init__(self): pass

    mod = make_module("mod7", {"Dummy6": Dummy6})
    docs, mode = extract_docstrings(mod) # 80.2μs -> 77.0μs (4.13% faster)

def test_extract_docstrings_class_with_inheritance():
    """Test extraction from a class that inherits from another."""

    class Base:
        """Base doc"""
        def __init__(self, x: int): """x: base param"""; pass

    class Sub(Base):
        """Sub doc"""
        def __init__(self, x: int, y: int): """x: sub x\ny: sub y"""; pass

    mod = make_module("mod8", {"Base": Base, "Sub": Sub})
    docs, mode = extract_docstrings(mod) # 10.2ms -> 193μs (5160% faster)




def test_extract_docstrings_many_classes_and_methods():
    """Test extraction from a module with many classes and methods."""
    # Create 50 classes, each with preprocess and postprocess
    classes = {}
    for i in range(50):
        def make_preprocess(i):
            def preprocess(self, x: int) -> int:
                """x: input x
                Returns: output x
                """
                return x + i
            return preprocess
        def make_postprocess(i):
            def postprocess(self, y: int) -> int:
                """y: input y
                Returns: output y
                """
                return y - i
            return postprocess
        cls = type(f"Class{i}", (), {
            "__doc__": f"Class{i} doc",
            "__init__": lambda self, x=0: None,
            "preprocess": make_preprocess(i),
            "postprocess": make_postprocess(i),
        })
        classes[f"Class{i}"] = cls
    mod = make_module("mod_large", classes)
    docs, mode = extract_docstrings(mod) # 334ms -> 5.65ms (5823% faster)
    # All classes should be present
    for i in range(50):
        cname = f"Class{i}"


def test_extract_docstrings_large_docstring_text():
    """Test extraction from a class with a very large docstring."""

    big_doc = "A" * 1000
    class BigDoc:
        __doc__ = big_doc
        def __init__(self): pass

    mod = make_module("mod_bigdoc", {"BigDoc": BigDoc})
    docs, mode = extract_docstrings(mod) # 80.0μs -> 73.2μs (9.22% faster)

def test_extract_docstrings_performance_many_classes():
    """Test extraction from a module with 100 classes, each with __init__."""

    classes = {}
    for i in range(100):
        def make_init(i):
            def __init__(self, x: int = i):
                """x: param x"""
                self.x = x
            return __init__
        cls = type(f"Class{i}", (), {
            "__doc__": f"Class{i} doc",
            "__init__": make_init(i),
        })
        classes[f"Class{i}"] = cls
    mod = make_module("mod_perf", classes)
    docs, mode = extract_docstrings(mod) # 654ms -> 4.89ms (13286% faster)
    # All classes should be present
    for i in range(100):
        cname = f"Class{i}"
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import sys
import textwrap
import types

# imports
import pytest
from gradio.cli.commands.components._docs_utils import extract_docstrings

# function to test
# (inserted verbatim from user prompt, so not repeated here)

# -----------------------------
# TEST MODULES AND CLASSES
# -----------------------------

# Helper: Create a dummy module dynamically for testing
def make_module(name, members):
    mod = types.ModuleType(name)
    for k, v in members.items():
        setattr(mod, k, v)
    sys.modules[name] = mod
    return mod

# Helper: Dummy event class for EVENTS testing
class DummyEvent:
    def __init__(self, doc):
        self.doc = doc

# 1. Basic Test Cases
def test_basic_function_and_class_docstring_extraction():
    # Define a simple class with __init__, preprocess, postprocess
    class Simple:
        """Simple class docstring.

        param1: Description of param1.
        Returns: Description of return.
        """
        def __init__(self, param1: int = 42):
            """Initializes Simple.

            param1: The first parameter.
            Returns: None.
            """
            pass

        def preprocess(self, value: str):
            """Preprocess value.

            value: The input string.
            Returns: None.
            """
            return value

        def postprocess(self, value: str) -> str:
            """Postprocess value.

            value: The processed string.
            Returns: The final string.
            """
            return value

    mod = make_module("mod_basic", {"Simple": Simple})
    docs, mode = extract_docstrings(mod) # 9.90ms -> 200μs (4827% faster)

def test_basic_function_with_no_docstring():
    # Function without docstring
    def func_no_doc(a: int):
        return a
    mod = make_module("mod_no_doc", {"func_no_doc": func_no_doc})
    docs, mode = extract_docstrings(mod) # 55.4μs -> 45.7μs (21.2% faster)

def test_basic_class_with_events():
    # Class with EVENTS attribute
    class WithEvents:
        """Class with events."""
        EVENTS = [DummyEvent("Event 1 for {{ component }}"), DummyEvent("Event 2 for {{ component }}")]
        def __init__(self):
            pass
    mod = make_module("mod_events", {"WithEvents": WithEvents})
    docs, mode = extract_docstrings(mod) # 76.4μs -> 71.9μs (6.28% faster)

# 2. Edge Test Cases

def test_parameter_with_no_type_hint_or_default():
    # __init__ param with no type hint or default
    class NoTypeOrDefault:
        """Class with param lacking type and default."""
        def __init__(self, x):
            """x: parameter x."""
            pass
    mod = make_module("mod_edge1", {"NoTypeOrDefault": NoTypeOrDefault})
    docs, mode = extract_docstrings(mod) # 66.8μs -> 63.2μs (5.77% faster)

def test_parameter_with_multiline_docstring():
    # Multiline param docstring
    class MultiLine:
        def __init__(self, foo: int = 1):
            """
            foo: This parameter
                spans multiple lines
                and should be captured.
            Returns: None.
            """
            pass
    mod = make_module("mod_edge2", {"MultiLine": MultiLine})
    docs, mode = extract_docstrings(mod) # 6.47ms -> 115μs (5503% faster)
    # Should capture only the first line after "foo:"
    desc = docs["MultiLine"]["members"]["__init__"]["foo"]["description"]

def test_function_with_unusual_parameter_names():
    # Parameter with underscores and digits
    def func_weird(_x1: float = 3.14):
        """_x1: A weird parameter."""
        return _x1
    mod = make_module("mod_edge3", {"func_weird": func_weird})
    docs, mode = extract_docstrings(mod) # 48.9μs -> 43.3μs (13.0% faster)

def test_class_with_no_members():
    # Class with no methods
    class Empty:
        """Empty class."""
        pass
    mod = make_module("mod_edge4", {"Empty": Empty})
    docs, mode = extract_docstrings(mod) # 43.0μs -> 38.8μs (10.6% faster)

def test_function_with_no_parameters():
    # Function with no parameters
    def func_no_params():
        """Returns: Always returns 1."""
        return 1
    mod = make_module("mod_edge5", {"func_no_params": func_no_params})
    docs, mode = extract_docstrings(mod) # 45.8μs -> 42.6μs (7.59% faster)

def test_class_with_inherited_method():
    class Base:
        def __init__(self, a: int): pass
    class Derived(Base):
        """Derived docstring."""
        def __init__(self, a: int, b: str = "hi"):
            """b: New param."""
            super().__init__(a)
    mod = make_module("mod_edge6", {"Derived": Derived})
    docs, mode = extract_docstrings(mod) # 9.82ms -> 129μs (7463% faster)

def test_class_with_multiple_methods():
    class MultiMethod:
        """MultiMethod doc."""
        def __init__(self, a: int): pass
        def preprocess(self, value: int): return value
        def postprocess(self, value: int): return value
        def some_other(self, x): return x
    mod = make_module("mod_edge7", {"MultiMethod": MultiMethod})
    docs, mode = extract_docstrings(mod) # 6.66ms -> 171μs (3792% faster)

def test_class_with_complex_type_hints():
    from typing import Dict, List, Union
    class ComplexType:
        def __init__(self, x: List[Dict[str, Union[int, float]]]):
            """x: Complex type."""
            pass
    mod = make_module("mod_edge8", {"ComplexType": ComplexType})
    docs, mode = extract_docstrings(mod) # 3.33ms -> 103μs (3119% faster)
    # Type should include List, Dict, Union, int, float, str
    typ = docs["ComplexType"]["members"]["__init__"]["x"]["type"]

def test_class_with_no_docstring_and_no_methods():
    class NoDocNoMeth:
        pass
    mod = make_module("mod_edge9", {"NoDocNoMeth": NoDocNoMeth})
    docs, mode = extract_docstrings(mod) # 43.6μs -> 38.1μs (14.3% faster)

# 3. Large Scale Test Cases

def test_large_module_many_classes_and_functions():
    # Create many classes and functions
    members = {}
    for i in range(50):
        # Each class has __init__, preprocess, postprocess
        class_name = f"Cls{i}"
        def make_class(idx):
            class Cls:
                """Class {idx} docstring."""
                def __init__(self, x: int = idx):
                    """x: parameter x."""
                    pass
                def preprocess(self, value: str):
                    """value: preprocess param."""
                    return value
                def postprocess(self, value: str):
                    """value: postprocess param."""
                    return value
            Cls.__name__ = class_name
            return Cls
        members[class_name] = make_class(i)
    for i in range(50):
        def func(x: int = i):
            """x: function param."""
            return x
        func.__name__ = f"func{i}"
        members[f"func{i}"] = func
    mod = make_module("mod_large", members)
    docs, mode = extract_docstrings(mod) # 494ms -> 5.92ms (8251% faster)
    # Check a few random classes and functions
    for idx in [0, 10, 25, 49]:
        cname = f"Cls{idx}"
        fname = f"func{idx}"
    # Should not be slow or crash

def test_large_module_with_long_docstrings():
    # Create a class with a long docstring
    long_doc = "This is a long docstring.\n" + "\n".join([f"Line {i}" for i in range(200)])
    class LongDocClass:
        """{doc}
        param1: The only parameter.
        Returns: None.
        """.format(doc=long_doc)
        def __init__(self, param1: int = 0):
            """param1: The only parameter."""
            pass
    mod = make_module("mod_longdoc", {"LongDocClass": LongDocClass})
    docs, mode = extract_docstrings(mod) # 6.44ms -> 110μs (5736% faster)
    # Should capture only the description without param lines
    desc = docs["LongDocClass"]["description"]

def test_large_module_with_many_events():
    # Class with many events
    events = [DummyEvent(f"Event {i} for {{ component }}") for i in range(100)]
    class ManyEvents:
        """Class with many events."""
        EVENTS = events
        def __init__(self): pass
    mod = make_module("mod_manyevents", {"ManyEvents": ManyEvents})
    docs, mode = extract_docstrings(mod) # 140μs -> 120μs (16.8% faster)
    for i in [0, 50, 99]:
        pass

def test_large_module_with_nested_classes():
    # Nested classes (should only extract top-level)
    class Outer:
        """Outer doc."""
        class Inner:
            """Inner doc."""
            def __init__(self, y: int): pass
        def __init__(self, x: int): pass
    mod = make_module("mod_nested", {"Outer": Outer})
    docs, mode = extract_docstrings(mod) # 3.28ms -> 106μs (2975% faster)

def test_large_module_with_varied_type_hints():
    # Classes with varied type hints
    from typing import Dict, List, Optional, Union
    class Varied:
        def __init__(self, a: Optional[int], b: List[str], c: Dict[str, Union[int, float]]):
            """a: optional int
            b: list of str
            c: dict of str to int or float"""
            pass
    mod = make_module("mod_varied", {"Varied": Varied})
    docs, mode = extract_docstrings(mod) # 9.79ms -> 135μs (7123% faster)
    typ_a = docs["Varied"]["members"]["__init__"]["a"]["type"]
    typ_b = docs["Varied"]["members"]["__init__"]["b"]["type"]
    typ_c = docs["Varied"]["members"]["__init__"]["c"]["type"]
# 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-extract_docstrings-mhwwzcx4 and push.