⚡️ Speed up function superreload by 99%

[codeflash-ai]

📄 99% (0.99x) speedup for superreload in marimo/_runtime/reload/autoreload.py

⏱️ Runtime : 9.75 seconds → 4.91 seconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 98% speedup through several key optimizations that reduce redundant operations and expensive attribute lookups:

Key Optimizations:

1. Eliminated Double Weakref Insertion: The original code redundantly called old_objects.setdefault(key, []).append(weakref.ref(obj)) both inside append_obj() and immediately after in the superreload() loop. The optimization removes this duplication, cutting the weakref creation overhead in half.

2. Reduced Attribute Lookups: Added local variables module_dict = module.__dict__ and mod_name = module.__name__ to cache frequently accessed attributes. The profiler shows these lookups were happening thousands of times per reload.

3. Optimized Dictionary Access Pattern: Replaced the key not in old_objects + old_objects[key] pattern with old_objects.get(key), eliminating redundant dictionary lookups.

4. Removed Unnecessary List Creation: Changed list(module.__dict__.items()) to module.__dict__.items() since the dictionary isn't modified during iteration, avoiding unnecessary list allocation.

5. Improved hasattr Pattern: Replaced hasattr(obj, "__module__") and obj.__module__ == module.__name__ with getattr(obj, "__module__", None) == module.__name__, which is more efficient and handles missing attributes gracefully.

Performance Impact by Test Scale:

• Small modules (single functions/classes): 0-4% improvement
• Medium modules (hundreds of objects): 20-25% improvement
• Large modules (500+ classes): 96-99% improvement

The optimizations are particularly effective for large-scale test cases because they eliminate O(n) redundant operations that compound significantly as module size increases. The update_generic function, which consumes 99% of runtime, benefits from receiving fewer duplicate calls due to the eliminated double-insertion.

✅ 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	90.7%

🌀 Generated Regression Tests and Runtime

import importlib
import os
import shutil
import sys
import tempfile
import types
import weakref

# imports
import pytest
from marimo._runtime.reload.autoreload import superreload


# --- Helper for creating and importing test modules dynamically ---
class TempModule:
    """Context manager to create a temporary Python module for testing."""

    def __init__(self, code: str, name: str = "testmod"):
        self.code = code
        self.name = name
        self.dir = None
        self.path = None

    def __enter__(self):
        self.dir = tempfile.mkdtemp()
        self.path = os.path.join(self.dir, f"{self.name}.py")
        with open(self.path, "w") as f:
            f.write(self.code)
        sys.path.insert(0, self.dir)
        importlib.invalidate_caches()
        self.module = importlib.import_module(self.name)
        return self

    def update_code(self, new_code: str):
        with open(self.path, "w") as f:
            f.write(new_code)
        importlib.invalidate_caches()

    def reload(self):
        return importlib.reload(self.module)

    def __exit__(self, exc_type, exc_val, exc_tb):
        sys.path.remove(self.dir)
        shutil.rmtree(self.dir)
        if self.name in sys.modules:
            del sys.modules[self.name]

# --- TESTS ---

# 1. Basic Test Cases

def test_superreload_function_updates_function_code():
    """Test that superreload updates the code object of a function."""
    code1 = "def foo():\n    return 1\n"
    code2 = "def foo():\n    return 2\n"
    with TempModule(code1, "mod1") as mod:
        old_objects = {}
        f = mod.module.foo
        mod.update_code(code2)
        superreload(mod.module, old_objects) # 153μs -> 152μs (0.432% faster)

def test_superreload_class_updates_class_dict():
    """Test that superreload updates the class dictionary for a class."""
    code1 = "class A:\n    def x(self):\n        return 1\n"
    code2 = "class A:\n    def x(self):\n        return 2\n"
    with TempModule(code1, "mod2") as mod:
        old_objects = {}
        A = mod.module.A
        a = A()
        mod.update_code(code2)
        superreload(mod.module, old_objects) # 8.08ms -> 4.12ms (96.2% faster)

def test_superreload_preserves_module_identity():
    """Test that superreload returns the same module object."""
    code = "x = 42\n"
    with TempModule(code, "mod3") as mod:
        old_objects = {}
        m1 = mod.module
        codeflash_output = superreload(m1, old_objects); m2 = codeflash_output # 129μs -> 124μs (3.75% faster)

def test_superreload_none_old_objects():
    """Test that superreload works when old_objects is None."""
    code = "def foo():\n    return 123\n"
    with TempModule(code, "mod4") as mod:
        mod.update_code("def foo():\n    return 456\n")
        codeflash_output = superreload(mod.module, None); m = codeflash_output # 152μs -> 147μs (3.98% faster)

def test_superreload_multiple_functions_and_classes():
    """Test that superreload updates multiple objects in the module."""
    code1 = (
        "def f(): return 1\n"
        "def g(): return 2\n"
        "class C:\n    def x(self): return 3\n"
    )
    code2 = (
        "def f(): return 10\n"
        "def g(): return 20\n"
        "class C:\n    def x(self): return 30\n"
    )
    with TempModule(code1, "mod5") as mod:
        old_objects = {}
        f, g, C = mod.module.f, mod.module.g, mod.module.C
        mod.update_code(code2)
        superreload(mod.module, old_objects) # 8.19ms -> 4.32ms (89.6% faster)

# 2. Edge Test Cases



def test_superreload_removes_deleted_objects():
    """Test that objects removed from module are no longer in old_objects."""
    code1 = "def foo(): return 1\n"
    code2 = ""
    with TempModule(code1, "mod9") as mod:
        old_objects = {}
        f = mod.module.foo
        mod.update_code(code2)
        superreload(mod.module, old_objects) # 199μs -> 200μs (0.426% slower)


def test_superreload_with_property():
    """Test that properties are handled (no-op in our stub)."""
    code1 = (
        "class A:\n"
        "    @property\n"
        "    def val(self):\n"
        "        return 1\n"
    )
    code2 = (
        "class A:\n"
        "    @property\n"
        "    def val(self):\n"
        "        return 2\n"
    )
    with TempModule(code1, "mod11") as mod:
        old_objects = {}
        A = mod.module.A
        a = A()
        mod.update_code(code2)
        superreload(mod.module, old_objects) # 8.13ms -> 4.15ms (96.0% faster)

def test_superreload_handles_module_with_no_loader():
    """Test that superreload doesn't crash if __loader__ is missing."""
    code = "x = 1\n"
    with TempModule(code, "mod12") as mod:
        del mod.module.__dict__["__loader__"]
        # Should not raise
        superreload(mod.module, {}) # 128μs -> 129μs (0.862% slower)

# 3. Large Scale Test Cases

def test_superreload_large_number_of_functions():
    """Test superreload with a large number of functions (scalability)."""
    n = 500
    code1 = "\n".join([f"def f{i}(): return {i}" for i in range(n)])
    code2 = "\n".join([f"def f{i}(): return {i*2}" for i in range(n)])
    with TempModule(code1, "mod13") as mod:
        old_objects = {}
        funcs = [getattr(mod.module, f"f{i}") for i in range(n)]
        for i, f in enumerate(funcs):
            pass
        mod.update_code(code2)
        superreload(mod.module, old_objects) # 3.63ms -> 3.02ms (20.0% faster)
        for i, f in enumerate(funcs):
            pass

def test_superreload_large_number_of_classes():
    """Test superreload with a large number of classes (scalability)."""
    n = 500
    code1 = "\n".join([f"class C{i}:\n    def val(self): return {i}" for i in range(n)])
    code2 = "\n".join([f"class C{i}:\n    def val(self): return {i*2}" for i in range(n)])
    with TempModule(code1, "mod14") as mod:
        old_objects = {}
        classes = [getattr(mod.module, f"C{i}") for i in range(n)]
        for i, C in enumerate(classes):
            pass
        mod.update_code(code2)
        superreload(mod.module, old_objects) # 4.09s -> 2.05s (99.0% faster)
        for i, C in enumerate(classes):
            pass

def test_superreload_large_mixed_objects():
    """Test superreload with a mix of many functions and classes."""
    n = 200
    code1 = (
        "\n".join([f"def f{i}(): return {i}" for i in range(n)]) + "\n" +
        "\n".join([f"class C{i}:\n    def val(self): return {i}" for i in range(n)])
    )
    code2 = (
        "\n".join([f"def f{i}(): return {i*3}" for i in range(n)]) + "\n" +
        "\n".join([f"class C{i}:\n    def val(self): return {i*4}" for i in range(n)])
    )
    with TempModule(code1, "mod15") as mod:
        old_objects = {}
        funcs = [getattr(mod.module, f"f{i}") for i in range(n)]
        classes = [getattr(mod.module, f"C{i}") for i in range(n)]
        for i in range(n):
            pass
        mod.update_code(code2)
        superreload(mod.module, old_objects) # 1.64s -> 823ms (99.3% faster)
        for i in range(n):
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import importlib
import os
import shutil
import sys
import tempfile
import types
import weakref

# imports
import pytest
from marimo._runtime.reload.autoreload import superreload

# --- Test utilities for dynamic module creation ---

def write_module(tmpdir, name, contents):
    path = os.path.join(tmpdir, name + ".py")
    with open(path, "w", encoding="utf-8") as f:
        f.write(contents)
    return path

@pytest.fixture
def tempmod(tmp_path):
    """Fixture to create a temporary module and clean up sys.modules"""
    created = []
    def _create(name, code):
        modpath = write_module(str(tmp_path), name, code)
        sys.path.insert(0, str(tmp_path))
        if name in sys.modules:
            del sys.modules[name]
        mod = importlib.import_module(name)
        created.append(name)
        return mod
    yield _create
    # cleanup
    for name in created:
        if name in sys.modules:
            del sys.modules[name]
    if str(tmp_path) in sys.path:
        sys.path.remove(str(tmp_path))

# --- Basic Test Cases ---

def test_superreload_function_update_basic(tempmod):
    """Test that superreload updates a function's code object in place."""
    code1 = "def foo():\n    return 1\n"
    mod = tempmod("mod1", code1)
    old_foo = mod.foo
    old_objects = {}
    # Change function body
    code2 = "def foo():\n    return 2\n"
    write_module(os.path.dirname(mod.__file__), "mod1", code2)
    codeflash_output = superreload(mod, old_objects); mod2 = codeflash_output # 154μs -> 157μs (1.50% slower)

def test_superreload_class_update_basic(tempmod):
    """Test that superreload updates a class's methods in place."""
    code1 = (
        "class A:\n"
        "    def foo(self): return 10\n"
    )
    mod = tempmod("mod2", code1)
    old_A = mod.A
    old_objects = {}
    code2 = (
        "class A:\n"
        "    def foo(self): return 20\n"
    )
    write_module(os.path.dirname(mod.__file__), "mod2", code2)
    codeflash_output = superreload(mod, old_objects); mod2 = codeflash_output # 8.29ms -> 4.21ms (96.7% faster)

def test_superreload_property_update_basic(tempmod):
    """Test that superreload updates a property in place."""
    code1 = (
        "class B:\n"
        "    @property\n"
        "    def x(self): return 1\n"
    )
    mod = tempmod("mod3", code1)
    old_B = mod.B
    old_objects = {}
    code2 = (
        "class B:\n"
        "    @property\n"
        "    def x(self): return 2\n"
    )
    write_module(os.path.dirname(mod.__file__), "mod3", code2)
    codeflash_output = superreload(mod, old_objects); mod2 = codeflash_output # 8.23ms -> 4.23ms (94.3% faster)

def test_superreload_preserves_non_module_objects(tempmod):
    """Test that objects not defined in the module are not tracked or updated."""
    code1 = (
        "import math\n"
        "def foo(): return math.sqrt(4)\n"
    )
    mod = tempmod("mod4", code1)
    old_objects = {}
    code2 = (
        "import math\n"
        "def foo(): return math.sqrt(9)\n"
    )
    write_module(os.path.dirname(mod.__file__), "mod4", code2)
    superreload(mod, old_objects) # 151μs -> 151μs (0.411% slower)
    # math should not be in old_objects
    for k in old_objects:
        pass

def test_superreload_returns_module(tempmod):
    """Test that superreload returns the reloaded module object."""
    code1 = "def foo(): return 1\n"
    mod = tempmod("mod5", code1)
    code2 = "def foo(): return 2\n"
    write_module(os.path.dirname(mod.__file__), "mod5", code2)
    codeflash_output = superreload(mod, {}); mod2 = codeflash_output # 143μs -> 139μs (3.03% faster)

# --- Edge Test Cases ---


def test_superreload_with_missing_module(tempmod):
    """Test that superreload raises when the module file is missing."""
    code1 = "def foo(): return 1\n"
    mod = tempmod("mod7", code1)
    os.remove(mod.__file__)
    with pytest.raises(ModuleNotFoundError):
        superreload(mod, {}) # 265μs -> 262μs (1.11% faster)

def test_superreload_with_non_module_object():
    """Test that superreload raises if passed a non-module object."""
    with pytest.raises(AttributeError):
        superreload(42, {}) # 1.70μs -> 1.53μs (11.2% faster)

def test_superreload_with_none_module():
    """Test that superreload raises if passed None as module."""
    with pytest.raises(AttributeError):
        superreload(None, {}) # 1.48μs -> 1.38μs (7.03% faster)


def test_superreload_with_object_without_module_attr(tempmod):
    """Test that append_obj returns False for objects without __module__."""
    class Dummy:
        pass
    mod = tempmod("mod8", "x = 1\n")
    d = {}

def test_superreload_with_object_weakref_fails(tempmod):
    """Test that append_obj skips objects that can't be weakref'd (e.g. int)."""
    mod = tempmod("mod9", "x = 1\n")
    d = {}

# --- Large Scale Test Cases ---

def test_superreload_large_module_many_functions(tempmod):
    """Test superreload on a module with many functions."""
    N = 500
    code1 = "\n".join([f"def f{i}(): return {i}" for i in range(N)])
    mod = tempmod("mod10", code1)
    old_funcs = [getattr(mod, f"f{i}") for i in range(N)]
    old_objects = {}
    # Change all return values
    code2 = "\n".join([f"def f{i}(): return {i+1}" for i in range(N)])
    write_module(os.path.dirname(mod.__file__), "mod10", code2)
    superreload(mod, old_objects) # 3.77ms -> 3.09ms (22.0% faster)
    # All old function objects should now return i+1
    for i, f in enumerate(old_funcs):
        pass

def test_superreload_large_module_many_classes(tempmod):
    """Test superreload on a module with many classes."""
    N = 300
    code1 = "\n".join([f"class C{i}:\n def foo(self): return {i}" for i in range(N)])
    mod = tempmod("mod11", code1)
    old_classes = [getattr(mod, f"C{i}") for i in range(N)]
    old_objects = {}
    code2 = "\n".join([f"class C{i}:\n def foo(self): return {i+2}" for i in range(N)])
    write_module(os.path.dirname(mod.__file__), "mod11", code2)
    superreload(mod, old_objects) # 2.38s -> 1.20s (98.1% faster)
    for i, cls in enumerate(old_classes):
        pass

def test_superreload_large_module_mixed(tempmod):
    """Test superreload on a module with many classes and functions."""
    N = 200
    code1 = "\n".join(
        [f"class C{i}:\n def foo(self): return {i}" for i in range(N)] +
        [f"def f{i}(): return {i*2}" for i in range(N)]
    )
    mod = tempmod("mod12", code1)
    old_classes = [getattr(mod, f"C{i}") for i in range(N)]
    old_funcs = [getattr(mod, f"f{i}") for i in range(N)]
    old_objects = {}
    code2 = "\n".join(
        [f"class C{i}:\n def foo(self): return {i+3}" for i in range(N)] +
        [f"def f{i}(): return {i*2+1}" for i in range(N)]
    )
    write_module(os.path.dirname(mod.__file__), "mod12", code2)
    superreload(mod, old_objects) # 1.59s -> 797ms (98.8% faster)
    for i, cls in enumerate(old_classes):
        pass
    for i, f in enumerate(old_funcs):
        pass

def test_superreload_large_module_performance(tempmod):
    """Test that superreload does not take excessive time for large modules."""
    import time
    N = 900
    code1 = "\n".join([f"def f{i}(): return {i}" for i in range(N)])
    mod = tempmod("mod13", code1)
    code2 = "\n".join([f"def f{i}(): return {i+10}" for i in range(N)])
    write_module(os.path.dirname(mod.__file__), "mod13", code2)
    start = time.time()
    superreload(mod, {}) # 6.45ms -> 5.41ms (19.3% faster)
    elapsed = time.time() - start
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from marimo._runtime.reload.autoreload import superreload

To edit these changes git checkout codeflash/optimize-superreload-mhb0qrp1 and push.