⚡️ Speed up function patch by 9%

[codeflash-ai]

📄 9% (0.09x) speedup for patch in wandb/integration/kfp/kfp_patch.py

⏱️ Runtime : 25.1 milliseconds → 23.1 milliseconds (best of 52 runs)

📝 Explanation and details

The optimized code achieves an 8% speedup through three key performance improvements:

1. Module Import Caching in get_module()

• Added _imported_modules cache to store successfully imported modules, avoiding expensive re-imports
• Early return for cached modules (line profiler shows 1010 hits with fast 206ns per hit)
• This is the most impactful optimization since import_module_lazy() dominates runtime (94.7% of total time)

2. Local Caching in full_path_exists()

• Eliminated the nested get_parent_child_pairs() function and inlined the logic
• Added module_cache dictionary to avoid repeated get_module() calls for the same parent modules within a single function call
• Reduced function call overhead and duplicate module lookups

3. String Processing Optimization in termerror()

• Added conditional check for newlines before expensive string splitting/joining operations
• For single-line error messages (the common case), uses simple string formatting instead of list comprehension
• Saves ~1.5ms on the string processing line (from 2.7ms to 0.7ms in line profiler)

The optimizations are particularly effective for the test workloads involving many module patches and repeated function calls, showing 2-74% improvements in individual test cases. The module caching provides the biggest benefit for scenarios with repeated imports of the same modules.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import sys
import types

# imports
import pytest
from wandb.integration.kfp.kfp_patch import patch

# --- Minimal mocks for wandb and wandb.util so we can test patch in isolation ---

class DummyWandb:
    def __init__(self):
        self.patched = {}
        self.termerror_calls = []

    def termerror(self, msg):
        self.termerror_calls.append(msg)

wandb = DummyWandb()
from wandb.integration.kfp.kfp_patch import patch

# --- Unit tests for patch ---

# Basic Test Cases









def test_patch_many_functions_in_one_module():
    """Patch a large number of functions in a single module."""
    module = types.ModuleType("largemod")
    sys.modules["largemod"] = module
    # Add 500 functions to module
    for i in range(500):
        setattr(module, f"func{i}", (lambda x=i: lambda: f"original_{x}")())
    # Patch all 500 functions
    patched_funcs = []
    for i in range(500):
        def make_func(x=i):
            def func(): return f"patched_{x}"
            func.__name__ = f"func{x}"
            return func
        func = make_func()
        codeflash_output = patch("largemod", func); result = codeflash_output # 2.65ms -> 1.92ms (37.5% faster)
        patched_funcs.append([module, f"func{i}"])

def test_patch_many_modules_one_func_each():
    """Patch the same-named function in many different modules."""
    for i in range(300):
        modname = f"mod_{i}"
        module = types.ModuleType(modname)
        sys.modules[modname] = module
        def func(): return f"original_{i}"
        func.__name__ = "func"
        setattr(module, "func", func)
        def func_patch(): return f"patched_{i}"
        func_patch.__name__ = "func"
        codeflash_output = patch(modname, func_patch); result = codeflash_output # 3.01ms -> 2.84ms (6.02% faster)

def test_patch_performance_large_scale():
    """Performance: Patch 1000 functions in 1000 modules (simulate, but only check a few)."""
    for i in range(0, 1000, 100):  # Only check every 100th to keep runtime reasonable
        modname = f"modp_{i}"
        module = types.ModuleType(modname)
        sys.modules[modname] = module
        def func(): return f"original_{i}"
        func.__name__ = f"func{i}"
        setattr(module, f"func{i}", func)
        def func_patch(): return f"patched_{i}"
        func_patch.__name__ = f"func{i}"
        codeflash_output = patch(modname, func_patch); result = codeflash_output # 35.7μs -> 20.5μs (74.2% faster)

# Edge Case: Patch on module with no __name__ attribute





#------------------------------------------------
import sys
import types

# imports
import pytest
from wandb.integration.kfp.kfp_patch import patch


# Mock wandb and its relevant submodules/classes for testing
class DummyWandbError(Exception):
    pass

class DummyTermError(Exception):
    def __init__(self, msg):
        self.msg = msg

class DummyWandbUtil:
    # Simulate a module registry
    _modules = {}

    @classmethod
    def get_module(cls, name, required=None, lazy=True):
        # Returns a dummy module if it exists, else None
        if name in cls._modules:
            return cls._modules[name]
        if required:
            raise DummyWandbError(required)
        return None

class DummyWandb:
    # patched is a dict mapping module names to list of [module, func_name] pairs
    patched = {}
    util = DummyWandbUtil

    @staticmethod
    def termerror(msg):
        # For testing, raise an exception so we can assert it was called
        raise DummyTermError(msg)

# Helper to create dummy modules with functions
def make_module(name, funcs):
    mod = types.ModuleType(name)
    for fname, fobj in funcs.items():
        setattr(mod, fname, fobj)
    mod.__name__ = name
    return mod
from wandb.integration.kfp.kfp_patch import patch

# Basic Test Cases

def test_patch_successful_basic():
    """Patch a function in a simple module and verify patching."""
    def orig_func(): return "original"
    def new_func(): return "patched"
    mod = make_module("mod1", {"func": orig_func})
    DummyWandbUtil._modules["mod1"] = mod
    DummyWandbUtil._modules["mod1.func"] = orig_func  # For full_path_exists

    # Patch
    codeflash_output = patch("mod1", new_func) # 31.9μs -> 31.0μs (2.91% faster)



def test_patch_module_with_nested_path():
    """Patch works with nested module paths."""
    def orig_func(): return "original"
    def new_func(): return "patched"
    # Simulate nested modules
    mod_parent = make_module("parent", {})
    mod_child = make_module("parent.child", {"func": orig_func})
    DummyWandbUtil._modules["parent"] = mod_parent
    DummyWandbUtil._modules["parent.child"] = mod_child
    DummyWandbUtil._modules["parent.child.func"] = orig_func

    codeflash_output = patch("parent.child", new_func) # 32.7μs -> 30.7μs (6.38% faster)

def test_patch_function_with_same_name_in_different_modules():
    """Patch functions with same name in different modules independently."""
    def orig_func1(): return "one"
    def orig_func2(): return "two"
    def new_func(): return "patched"
    mod1 = make_module("modA", {"func": orig_func1})
    mod2 = make_module("modB", {"func": orig_func2})
    DummyWandbUtil._modules["modA"] = mod1
    DummyWandbUtil._modules["modA.func"] = orig_func1
    DummyWandbUtil._modules["modB"] = mod2
    DummyWandbUtil._modules["modB.func"] = orig_func2

    codeflash_output = patch("modA", new_func) # 31.5μs -> 30.3μs (3.78% faster)
    codeflash_output = patch("modB", new_func) # 17.8μs -> 17.4μs (2.33% faster)

def test_patch_function_with_nonstandard_name():
    """Patch a function with underscores and numbers in its name."""
    def orig_func_42(): return "original"
    def new_func_42(): return "patched"
    mod = make_module("modX", {"func_42": orig_func_42})
    DummyWandbUtil._modules["modX"] = mod
    DummyWandbUtil._modules["modX.func_42"] = orig_func_42

    # The patch function's __name__ must match
    new_func_42.__name__ = "func_42"
    codeflash_output = patch("modX", new_func_42) # 31.2μs -> 29.6μs (5.68% faster)

# Large Scale Test Cases

def test_patch_many_functions_in_one_module():
    """Patch a large number of functions in a single module."""
    N = 500  # Large but <1000
    orig_funcs = {f"f{i}": (lambda i=i: f"orig{i}") for i in range(N)}
    mod = make_module("bigmod", orig_funcs)
    for fname, fobj in orig_funcs.items():
        DummyWandbUtil._modules[f"bigmod.{fname}"] = fobj
    DummyWandbUtil._modules["bigmod"] = mod

    # Patch all functions
    for i in range(N):
        def new_func(i=i): return f"patch{i}"
        new_func.__name__ = f"f{i}"
        codeflash_output = patch("bigmod", new_func) # 4.70ms -> 4.53ms (3.71% faster)

def test_patch_many_modules_one_function_each():
    """Patch one function in many modules."""
    N = 500
    for i in range(N):
        def orig_func(i=i): return f"orig{i}"
        def new_func(i=i): return f"patch{i}"
        modname = f"mod_{i}"
        mod = make_module(modname, {"func": orig_func})
        DummyWandbUtil._modules[modname] = mod
        DummyWandbUtil._modules[f"{modname}.func"] = orig_func
        codeflash_output = patch(modname, new_func) # 5.09ms -> 4.65ms (9.56% faster)

def test_patch_performance_large_scale(monkeypatch):
    """Performance: patching 1000 functions should not be slow."""
    import time
    N = 1000
    orig_funcs = {f"f{i}": (lambda i=i: f"orig{i}") for i in range(N)}
    mod = make_module("perfmod", orig_funcs)
    for fname, fobj in orig_funcs.items():
        DummyWandbUtil._modules[f"perfmod.{fname}"] = fobj
    DummyWandbUtil._modules["perfmod"] = mod

    start = time.time()
    for i in range(N):
        def new_func(i=i): return f"patch{i}"
        new_func.__name__ = f"f{i}"
        codeflash_output = patch("perfmod", new_func) # 9.32ms -> 8.82ms (5.69% faster)
    elapsed = time.time() - start

# Additional edge case: patch function with same name as builtin
def test_patch_function_named_builtin():
    """Patch a function named 'print'."""
    def orig_print(): return "original"
    def new_print(): return "patched"
    mod = make_module("mod_print", {"print": orig_print})
    DummyWandbUtil._modules["mod_print"] = mod
    DummyWandbUtil._modules["mod_print.print"] = orig_print
    new_print.__name__ = "print"
    codeflash_output = patch("mod_print", new_print) # 31.4μs -> 30.5μs (3.07% faster)

# Additional edge case: patch function with unicode name (Python 3 only)
def test_patch_function_with_unicode_name():
    """Patch a function with a unicode name."""
    def orig_func(): return "original"
    def new_func(): return "patched"
    # Give the function a unicode name
    new_func.__name__ = "func_üñîçødë"
    mod = make_module("mod_unicode", {"func_üñîçødë": orig_func})
    DummyWandbUtil._modules["mod_unicode"] = mod
    DummyWandbUtil._modules["mod_unicode.func_üñîçødë"] = orig_func
    codeflash_output = patch("mod_unicode", new_func) # 32.8μs -> 32.1μs (2.07% 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-patch-mhdm1kis and push.