From a5ec8998c504873c3a4b9bb565495442c30f87fe Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Thu, 13 Nov 2025 09:18:57 +0000
Subject: [PATCH] Optimize format_docstring
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 mlflow/utils/docstring_utils.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/mlflow/utils/docstring_utils.py b/mlflow/utils/docstring_utils.py
index 6b957beaa9a03..59a107582988c 100644
--- a/mlflow/utils/docstring_utils.py
+++ b/mlflow/utils/docstring_utils.py
@@ -152,10 +152,15 @@ def format_docstring(param_docs):
             p2: doc2
                 doc2 second line
     """
-    param_docs = ParamDocs(param_docs)
+    # Avoid double wrapping if already ParamDocs instance.
+    if type(param_docs) is not ParamDocs:
+        param_docs = ParamDocs(param_docs)
+
 
     def decorator(func):
-        func.__doc__ = param_docs.format_docstring(func.__doc__)
+        doc = func.__doc__
+        if doc is not None:
+            func.__doc__ = param_docs.format_docstring(doc)
         return func
 
     return decorator