⚡️ Speed up function format_docstring by 5%

[codeflash-ai]

📄 5% (0.05x) speedup for format_docstring in mlflow/utils/docstring_utils.py

⏱️ Runtime : 47.8 microseconds → 45.4 microseconds (best of 78 runs)

📝 Explanation and details

The optimized code applies two key performance optimizations:

1. Avoid redundant ParamDocs construction: The original code unconditionally wraps param_docs in ParamDocs(), even if it's already a ParamDocs instance. The optimization adds a type check if type(param_docs) is not ParamDocs: to only construct a new instance when needed. Using type() is not is faster than isinstance() for single-type checks since it avoids method dispatch overhead.

2. Skip processing for None docstrings: The optimization adds a null check if doc is not None: before calling format_docstring(). This avoids unnecessary string processing when functions have no docstring, which is common in Python codebases.

Performance impact: These optimizations show a 5% speedup (47.8μs → 45.4μs). The line profiler reveals that while the type check adds some overhead (35.2% of time), it prevents the more expensive ParamDocs() construction in cases where it's not needed. The null check for docstrings is essentially free but saves processing time when applicable.

Test case benefits: The optimizations particularly benefit scenarios with:

• Large-scale operations (100+ parameters): 5.21% faster in test_many_params_replacement
• Functions with no docstrings: 1.19% faster in test_none_docstring
• Large docstrings with no placeholders: 11.2% faster in test_large_docstring_no_placeholders

This decorator is likely used extensively in MLflow's API documentation generation, making these micro-optimizations valuable for reducing overall import and initialization time across the codebase.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import textwrap

# imports
import pytest
from mlflow.utils.docstring_utils import format_docstring


# Supporting class as implied by the function (minimal stub)
class ParamDocs:
    """
    Minimal implementation of ParamDocs to support format_docstring's usage.
    It expects a dict-like input and replaces {{ param }} placeholders in docstrings.
    """
    def __init__(self, docs):
        if isinstance(docs, ParamDocs):
            self.docs = docs.docs.copy()
        elif isinstance(docs, dict):
            self.docs = docs.copy()
        else:
            raise TypeError("ParamDocs expects a dict or ParamDocs instance")
    
    def format_docstring(self, docstring):
        if docstring is None:
            return None
        result = docstring
        for key, value in self.docs.items():
            # Replace all occurrences of {{ key }} with value (preserve indentation for multi-line)
            placeholder = "{{ %s }}" % key
            if placeholder not in result:
                continue
            # Handle multi-line value indentation
            lines = result.splitlines()
            for i, line in enumerate(lines):
                idx = line.find(placeholder)
                if idx != -1:
                    # Get indentation
                    indent = line[:idx]
                    # Split value into lines and indent accordingly
                    value_lines = value.splitlines()
                    value_lines = [value_lines[0]] + [
                        indent + vl if vl.strip() else vl for vl in value_lines[1:]
                    ]
                    lines[i] = line.replace(placeholder, value_lines[0], 1)
                    # Insert following lines if needed
                    if len(value_lines) > 1:
                        for offset, vline in enumerate(value_lines[1:], 1):
                            lines.insert(i + offset, vline)
            result = "\n".join(lines)
        return result
from mlflow.utils.docstring_utils import format_docstring

# unit tests

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

def test_single_param_replacement():
    """Test basic replacement of a single parameter."""
    param_docs = {"p1": "This is parameter one."}
    @format_docstring(param_docs)
    def func(p1):
        """
        Args:
            p1: {{ p1 }}
        """
    expected = """
    Args:
        p1: This is parameter one.
    """

def test_multiple_param_replacement():
    """Test replacement of multiple parameters."""
    param_docs = {"p1": "Doc for p1.", "p2": "Doc for p2."}
    @format_docstring(param_docs)
    def func(p1, p2):
        """
        Args:
            p1: {{ p1 }}
            p2: {{ p2 }}
        """
    expected = """
    Args:
        p1: Doc for p1.
        p2: Doc for p2.
    """

def test_multiline_param_doc():
    """Test that multi-line parameter documentation is indented properly."""
    param_docs = {"p1": "First line.\nSecond line."}
    @format_docstring(param_docs)
    def func(p1):
        """
        Args:
            p1: {{ p1 }}
        """
    expected = """
    Args:
        p1: First line.
            Second line.
    """

def test_no_placeholder():
    """Test that docstring is unchanged if there are no placeholders."""
    param_docs = {"p1": "Doc for p1."}
    @format_docstring(param_docs)
    def func(p1):
        """
        Args:
            p2: No placeholder here.
        """
    expected = """
    Args:
        p2: No placeholder here.
    """

def test_unused_param_docs():
    """Test that extra param_docs keys do not affect the docstring."""
    param_docs = {"p1": "Doc for p1.", "unused": "Should not appear"}
    @format_docstring(param_docs)
    def func(p1):
        """
        Args:
            p1: {{ p1 }}
        """
    expected = """
    Args:
        p1: Doc for p1.
    """

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

def test_empty_docstring():
    """Test when the function has an empty docstring."""
    param_docs = {"p1": "Doc for p1."}
    @format_docstring(param_docs)
    def func(p1):
        ""

def test_none_docstring():
    """Test when the function has no docstring (None)."""
    param_docs = {"p1": "Doc for p1."}
    @format_docstring(param_docs)
    def func(p1):
        pass

def test_placeholder_missing_in_param_docs():
    """Test when the docstring has a placeholder not present in param_docs."""
    param_docs = {"p1": "Doc for p1."}
    @format_docstring(param_docs)
    def func(p1, p2):
        """
        Args:
            p1: {{ p1 }}
            p2: {{ p2 }}
        """
    expected = """
    Args:
        p1: Doc for p1.
        p2: {{ p2 }}
    """

def test_placeholder_with_extra_spaces():
    """Test that placeholders with extra spaces are not replaced."""
    param_docs = {"p1": "Doc for p1."}
    @format_docstring(param_docs)
    def func(p1):
        """
        Args:
            p1: {{  p1  }}
        """
    expected = """
    Args:
        p1: {{  p1  }}
    """


def test_placeholder_at_start_of_line():
    """Test that placeholder at the start of a line is replaced correctly."""
    param_docs = {"desc": "Description here."}
    @format_docstring(param_docs)
    def func():
        """
        {{ desc }}
        """
    expected = """
    Description here.
    """

def test_placeholder_multiple_times():
    """Test that multiple occurrences of the same placeholder are all replaced."""
    param_docs = {"p1": "Doc for p1."}
    @format_docstring(param_docs)
    def func():
        """
        {{ p1 }}
        Again: {{ p1 }}
        """
    expected = """
    Doc for p1.
    Again: Doc for p1.
    """

def test_placeholder_with_similar_names():
    """Test that placeholders with similar names are not confused."""
    param_docs = {"p": "Doc for p.", "p1": "Doc for p1."}
    @format_docstring(param_docs)
    def func():
        """
        p: {{ p }}
        p1: {{ p1 }}
        """
    expected = """
    p: Doc for p.
    p1: Doc for p1.
    """

def test_placeholder_with_empty_param_doc():
    """Test that empty string param doc replaces placeholder with nothing."""
    param_docs = {"p1": ""}
    @format_docstring(param_docs)
    def func():
        """
        Args:
            p1: {{ p1 }}
        """
    expected = """
    Args:
        p1: 
    """


def test_large_number_of_params():
    """Test with a large number of parameters and placeholders."""
    param_docs = {f"p{i}": f"Doc for p{i}." for i in range(100)}
    docstring = "Args:\n" + "\n".join([f"    p{i}: {{ {{ p{i} }} }}" for i in range(100)])
    expected = "Args:\n" + "\n".join([f"    p{i}: Doc for p{i}." for i in range(100)])
    @format_docstring(param_docs)
    def func(*args):
        """
        {}
        """.format(docstring)

def test_large_multiline_param_docs():
    """Test with many multi-line parameter docs."""
    param_docs = {f"p{i}": f"Doc{i} line1\nDoc{i} line2" for i in range(50)}
    docstring = "Args:\n" + "\n".join([f"    p{i}: {{ {{ p{i} }} }}" for i in range(50)])
    expected_lines = ["Args:"]
    for i in range(50):
        expected_lines.append(f"    p{i}: Doc{i} line1")
        expected_lines.append(f"        Doc{i} line2")
    expected = "\n".join(expected_lines)
    @format_docstring(param_docs)
    def func(*args):
        """
        {}
        """.format(docstring)

def test_large_docstring_with_no_placeholders():
    """Test that a large docstring with no placeholders is unchanged."""
    param_docs = {f"p{i}": f"Doc for p{i}." for i in range(100)}
    docstring = "Args:\n" + "\n".join([f"    p{i}: no placeholder" for i in range(100)])
    @format_docstring(param_docs)
    def func(*args):
        """
        {}
        """.format(docstring)

import textwrap

# imports
import pytest
from mlflow.utils.docstring_utils import format_docstring


# Minimal ParamDocs implementation for testing
class ParamDocs:
    """
    Utility class to format docstrings by replacing placeholders with parameter documentation.
    """
    def __init__(self, param_docs):
        # Accept dict or ParamDocs (copy internal dict if so)
        if isinstance(param_docs, ParamDocs):
            self.param_docs = dict(param_docs.param_docs)
        elif isinstance(param_docs, dict):
            self.param_docs = dict(param_docs)
        else:
            raise TypeError("param_docs must be a dict or ParamDocs instance")

    def format_docstring(self, docstring):
        """
        Replace all {{ param_name }} occurrences with their documentation.
        Handles multiline docs and preserves indentation.
        """
        if docstring is None:
            return None
        lines = docstring.split('\n')
        formatted_lines = []
        for line in lines:
            # Find all placeholders in the line
            start = 0
            while True:
                idx1 = line.find('{{', start)
                if idx1 == -1:
                    break
                idx2 = line.find('}}', idx1)
                if idx2 == -1:
                    break
                key = line[idx1+2:idx2].strip()
                if key not in self.param_docs:
                    raise KeyError(f"Parameter '{key}' not found in param_docs")
                doc = self.param_docs[key]
                # If doc is multiline, indent subsequent lines
                doc_lines = doc.split('\n')
                indent = ' ' * (len(line) - len(line.lstrip()))
                # Replace placeholder with first line of doc
                line = line[:idx1] + doc_lines[0] + line[idx2+2:]
                # If multiline, add extra lines with indentation
                for extra_line in doc_lines[1:]:
                    formatted_lines.append(indent + extra_line)
                start = idx1 + len(doc_lines[0])
            formatted_lines.append(line)
        return '\n'.join(formatted_lines)
from mlflow.utils.docstring_utils import format_docstring

# unit tests

# 1. Basic Test Cases

def test_single_param_replacement():
    """Test single placeholder replacement."""
    param_docs = {"p1": "This is param p1."}
    @format_docstring(param_docs)
    def func(p1):
        """Args:
            p1: {{ p1 }}
        """
    expected = "Args:\n    p1: This is param p1."

def test_multiple_param_replacement():
    """Test multiple placeholder replacements."""
    param_docs = {"p1": "Doc1", "p2": "Doc2"}
    @format_docstring(param_docs)
    def func(p1, p2):
        """Args:
            p1: {{ p1 }}
            p2: {{ p2 }}
        """
    expected = "Args:\n    p1: Doc1\n    p2: Doc2"

def test_multiline_param_doc():
    """Test multiline param doc replacement and indentation."""
    param_docs = {"p1": "Doc1\nSecond line"}
    @format_docstring(param_docs)
    def func(p1):
        """Args:
            p1: {{ p1 }}
        """
    expected = "Args:\n    p1: Doc1\n        Second line"

def test_no_placeholder():
    """Test docstring with no placeholders remains unchanged."""
    param_docs = {"p1": "Doc1"}
    @format_docstring(param_docs)
    def func():
        """No params here."""
    expected = "No params here."



def test_empty_docstring():
    """Test empty docstring returns empty string."""
    param_docs = {"p1": "Doc1"}
    @format_docstring(param_docs)
    def func(p1):
        """"""
    expected = ""

def test_none_docstring():
    """Test None docstring returns None."""
    param_docs = {"p1": "Doc1"}
    def func(p1):
        pass
    func.__doc__ = None
    decorated = format_docstring(param_docs)(func) # 1.19μs -> 1.17μs (1.19% faster)

def test_placeholder_with_extra_spaces():
    """Test placeholder with extra spaces is handled."""
    param_docs = {"p1": "Doc1"}
    @format_docstring(param_docs)
    def func(p1):
        """Args:
            p1: {{    p1    }}
        """
    expected = "Args:\n    p1: Doc1"

def test_placeholder_at_start_of_line():
    """Test placeholder at start of line with multiline doc."""
    param_docs = {"p1": "Doc1\nSecond"}
    @format_docstring(param_docs)
    def func(p1):
        """{{ p1 }}"""
    expected = "Doc1\nSecond"

def test_multiple_placeholders_in_one_line():
    """Test multiple placeholders in a single line."""
    param_docs = {"p1": "Doc1", "p2": "Doc2"}
    @format_docstring(param_docs)
    def func(p1, p2):
        """Args: {{ p1 }}, {{ p2 }}"""
    expected = "Args: Doc1, Doc2"

def test_placeholder_with_no_surrounding_spaces():
    """Test placeholder with no spaces between braces and key."""
    param_docs = {"p1": "Doc1"}
    @format_docstring(param_docs)
    def func(p1):
        """Args:
            p1: {{p1}}
        """
    expected = "Args:\n    p1: Doc1"

def test_placeholder_with_special_characters_in_param_name():
    """Test param name with underscores and digits."""
    param_docs = {"p_1a": "Doc1"}
    @format_docstring(param_docs)
    def func(p_1a):
        """Args:
            p_1a: {{ p_1a }}
        """
    expected = "Args:\n    p_1a: Doc1"


def test_docstring_with_only_placeholder():
    """Test docstring is only a placeholder."""
    param_docs = {"p1": "Doc1"}
    @format_docstring(param_docs)
    def func(p1):
        """{{ p1 }}"""
    expected = "Doc1"

def test_docstring_with_placeholder_and_text():
    """Test placeholder surrounded by text."""
    param_docs = {"p1": "Doc1"}
    @format_docstring(param_docs)
    def func(p1):
        """Start {{ p1 }} end"""
    expected = "Start Doc1 end"

def test_placeholder_with_empty_docstring_value():
    """Test param_docs with empty string value."""
    param_docs = {"p1": ""}
    @format_docstring(param_docs)
    def func(p1):
        """Args:
            p1: {{ p1 }}
        """
    expected = "Args:\n    p1: "

# 3. Large Scale Test Cases

def test_many_params_replacement():
    """Test replacement with many parameters (scalability)."""
    param_docs = {f"p{i}": f"Doc{i}" for i in range(100)}
    doc_lines = ["Args:"]
    for i in range(100):
        doc_lines.append(f"    p{i}: {{ p{i} }}")
    docstring = "\n".join(doc_lines)
    # Dynamically create function with large docstring
    def func(*args):
        pass
    func.__doc__ = docstring
    decorated = format_docstring(param_docs)(func) # 1.61μs -> 1.53μs (5.21% faster)
    expected_lines = ["Args:"]
    for i in range(100):
        expected_lines.append(f"    p{i}: Doc{i}")
    expected = "\n".join(expected_lines)

def test_many_multiline_params():
    """Test replacement with many multiline parameters."""
    param_docs = {f"p{i}": f"Doc{i}\nLine2-{i}" for i in range(50)}
    doc_lines = ["Args:"]
    for i in range(50):
        doc_lines.append(f"    p{i}: {{ p{i} }}")
    docstring = "\n".join(doc_lines)
    def func(*args):
        pass
    func.__doc__ = docstring
    decorated = format_docstring(param_docs)(func) # 1.50μs -> 1.45μs (3.74% faster)
    expected_lines = ["Args:"]
    for i in range(50):
        expected_lines.append(f"    p{i}: Doc{i}")
        expected_lines.append(f"        Line2-{i}")
    expected = "\n".join(expected_lines)

def test_large_docstring_with_few_placeholders():
    """Test large docstring with few placeholders."""
    param_docs = {"p1": "Doc1", "p2": "Doc2"}
    doc_lines = ["Header"]
    for i in range(500):
        doc_lines.append(f"Line {i}")
    doc_lines.append("p1: {{ p1 }}")
    doc_lines.append("p2: {{ p2 }}")
    docstring = "\n".join(doc_lines)
    def func():
        pass
    func.__doc__ = docstring
    decorated = format_docstring(param_docs)(func) # 1.19μs -> 1.11μs (6.82% faster)
    expected_lines = ["Header"]
    for i in range(500):
        expected_lines.append(f"Line {i}")
    expected_lines.append("p1: Doc1")
    expected_lines.append("p2: Doc2")
    expected = "\n".join(expected_lines)

def test_large_param_docs_with_small_docstring():
    """Test large param_docs dict with small docstring."""
    param_docs = {f"p{i}": f"Doc{i}" for i in range(1000)}
    @format_docstring(param_docs)
    def func(p1):
        """Args:
            p1: {{ p1 }}
        """
    expected = "Args:\n    p1: Doc1"

def test_large_docstring_no_placeholders():
    """Test large docstring with no placeholders."""
    param_docs = {"p1": "Doc1"}
    doc_lines = [f"Line {i}" for i in range(1000)]
    docstring = "\n".join(doc_lines)
    def func():
        pass
    func.__doc__ = docstring
    decorated = format_docstring(param_docs)(func) # 1.26μs -> 1.13μs (11.2% faster)
    expected = "\n".join(doc_lines)
# 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-format_docstring-mhx7vs5r and push.