⚡️ Speed up function superreload by 101%

[codeflash-ai]

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

⏱️ Runtime : 10.7 seconds → 5.31 seconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 100% speedup (2x faster) through several key optimizations in the superreload function and related helpers:

Key Optimizations

1. Eliminated Duplicate Weakref Creation

• Original: Both append_obj and superreload created weakrefs for the same objects, doing duplicate work in the hot path
• Optimized: Removed redundant weakref creation in superreload since append_obj already handles this
• Impact: Reduces ~50% of weakref.ref() calls during object collection phase

2. Improved Object Module Check

• Original: hasattr(obj, "__module__") and obj.__module__ == module.__name__ (2 attribute lookups)
• Optimized: getattr(obj, "__module__", None) != module.__name__ (1 attribute lookup with default)
• Impact: 23% faster in append_obj function, reducing overhead for each object check

3. Eliminated Unnecessary list() Conversion

• Original: for name, obj in list(module.__dict__.items()):
• Optimized: for name, obj in module.__dict__.items():
• Impact: Avoids creating temporary lists, reducing memory allocation and iteration overhead

4. Optimized Dictionary Access Pattern

• Original: if key not in old_objects: followed by old_objects[key]
• Optimized: old_refs = old_objects.get(key) with None check
• Impact: Single dictionary lookup instead of two, reducing hash table operations

5. Enhanced Error Handling

• Original: old_dict["__loader__"] (potential KeyError)
• Optimized: old_dict.get("__loader__") (safe access)
• Impact: Avoids exception handling overhead in edge cases

6. Early Return in write_traceback

• Original: Always checked isinstance, then branched
• Optimized: Short-circuit with early return for common case
• Impact: Reduces function call overhead when not using custom stderr

Performance Impact by Test Case

The optimizations show consistent improvements across all test scenarios:

• Simple cases: 3-8% faster (basic function/class reloads)
• Complex cases: 95-100% faster (large modules with many classes/functions)
• Large-scale tests: Up to 102% faster for modules with 900+ objects

The optimizations are particularly effective for large modules and frequent reloading scenarios where the reduced object processing overhead and eliminated duplicate work compound significantly. Based on the test results, workloads with many classes and complex object hierarchies benefit most from these optimizations.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 24 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

--- Pytest Fixtures and Helpers ---

@pytest.fixture
def temp_module_dir():
# Create a temporary directory for our test modules
dirpath = tempfile.mkdtemp()
sys.path.insert(0, dirpath)
yield dirpath
sys.path.remove(dirpath)
shutil.rmtree(dirpath)

def write_module(dirpath, modname, code):
# Write a Python module to the given directory
filepath = os.path.join(dirpath, modname + ".py")
with open(filepath, "w") as f:
f.write(code)
return filepath

def import_module(modname):
# Import a module freshly (remove from sys.modules first)
if modname in sys.modules:
del sys.modules[modname]
return importlib.import_module(modname)

--- Basic Test Cases ---

def test_reload_simple_function(temp_module_dir):
"""
Test superreload upgrades a function's code object after module change.
"""
modname = "mod_simple"
code1 = "def foo():\n return 1\n"
code2 = "def foo():\n return 2\n"

write_module(temp_module_dir, modname, code1)
mod = import_module(modname)
old_objects = {}
f = mod.foo

# Update module code
write_module(temp_module_dir, modname, code2)
superreload(mod, old_objects) # 152μs -> 147μs (3.62% faster)

def test_reload_class_method(temp_module_dir):
"""
Test superreload upgrades a class method after module change.
"""
modname = "mod_class"
code1 = "class C:\n def foo(self):\n return 10\n"
code2 = "class C:\n def foo(self):\n return 20\n"

write_module(temp_module_dir, modname, code1)
mod = import_module(modname)
old_objects = {}
inst = mod.C()

# Update module code
write_module(temp_module_dir, modname, code2)
superreload(mod, old_objects) # 6.97ms -> 3.57ms (95.3% faster)

def test_reload_property(temp_module_dir):
"""
Test superreload upgrades a property after module change.
"""
modname = "mod_property"
code1 = ("class C:\n"
" @Property\n"
" def foo(self):\n"
" return 'a'\n")
code2 = ("class C:\n"
" @Property\n"
" def foo(self):\n"
" return 'b'\n")

write_module(temp_module_dir, modname, code1)
mod = import_module(modname)
old_objects = {}
inst = mod.C()

write_module(temp_module_dir, modname, code2)
superreload(mod, old_objects) # 7.00ms -> 3.59ms (95.0% faster)

def test_reload_multiple_objects(temp_module_dir):
"""
Test superreload upgrades multiple objects (function, class, property).
"""
modname = "mod_multi"
code1 = (
"def foo(): return 1\n"
"class C:\n"
" @Property\n"
" def bar(self): return 2\n"
)
code2 = (
"def foo(): return 3\n"
"class C:\n"
" @Property\n"
" def bar(self): return 4\n"
)
write_module(temp_module_dir, modname, code1)
mod = import_module(modname)
old_objects = {}
f = mod.foo
inst = mod.C()

write_module(temp_module_dir, modname, code2)
superreload(mod, old_objects) # 7.06ms -> 3.57ms (97.5% faster)

--- Edge Test Cases ---

def test_reload_non_module_object():
"""
Test superreload raises if passed a non-module object.
"""
with pytest.raises(AttributeError):
superreload(object(), {}) # 1.52μs -> 1.66μs (8.44% slower)

def test_reload_module_with_no_loader(temp_module_dir):
"""
Test superreload on module missing loader.
"""
modname = "mod_noloader"
code = "def foo(): return 1\n"
write_module(temp_module_dir, modname, code)
mod = import_module(modname)
old_objects = {}

# Remove __loader__
del mod.__dict__["__loader__"]
# Should not raise
superreload(mod, old_objects) # 149μs -> 142μs (5.36% faster)

def test_reload_object_deleted(temp_module_dir):
"""
Test superreload cleans up old_objects mapping if object is deleted.
"""
modname = "mod_del"
code1 = "def foo(): return 1\n"
code2 = "def foo(): return 2\n"

write_module(temp_module_dir, modname, code1)
mod = import_module(modname)
old_objects = {}
f = mod.foo

# Remove reference
del f
import gc; gc.collect()

write_module(temp_module_dir, modname, code2)
superreload(mod, old_objects) # 212μs -> 202μs (4.73% faster)
# old_objects mapping for foo should be cleaned up
foo_key = (modname, "foo")

def test_reload_module_with_unreferencable_object(temp_module_dir):
"""
Test superreload skips objects that cannot be weakref'd (like int).
"""
modname = "mod_int"
code = "x = 123\n"
write_module(temp_module_dir, modname, code)
mod = import_module(modname)
old_objects = {}
# Should not raise
superreload(mod, old_objects) # 136μs -> 131μs (3.67% faster)
# old_objects mapping for x should exist but be empty
x_key = (modname, "x")

--- Large Scale Test Cases ---

def test_reload_many_functions(temp_module_dir):
"""
Test superreload upgrades many functions efficiently.
"""
modname = "mod_manyfunc"
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+1}" for i in range(N)])

write_module(temp_module_dir, modname, code1)
mod = import_module(modname)
old_objects = {}
funcs = [getattr(mod, f"f{i}") for i in range(N)]

write_module(temp_module_dir, modname, code2)
superreload(mod, old_objects) # 3.82ms -> 3.17ms (20.2% faster)

def test_reload_many_classes(temp_module_dir):
"""
Test superreload upgrades many classes efficiently.
"""
modname = "mod_manyclass"
N = 300
code1 = "\n".join([f"class C{i}:\n def foo(self): return {i}" for i in range(N)])
code2 = "\n".join([f"class C{i}:\n def foo(self): return {i+1}" for i in range(N)])

write_module(temp_module_dir, modname, code1)
mod = import_module(modname)
old_objects = {}
instances = [getattr(mod, f"C{i}")() for i in range(N)]

write_module(temp_module_dir, modname, code2)
superreload(mod, old_objects) # 2.06s -> 1.03s (100% faster)

def test_reload_large_module_namespace(temp_module_dir):
"""
Test superreload handles modules with large namespace.
"""
modname = "mod_large"
N = 900
code1 = "\n".join([
f"def f{i}(): return {i}\nclass C{i}: pass\nx{i} = {i}"
for i in range(N)
])
code2 = "\n".join([
f"def f{i}(): return {i+2}\nclass C{i}: pass\nx{i} = {i+2}"
for i in range(N)
])
write_module(temp_module_dir, modname, code1)
mod = import_module(modname)
old_objects = {}
funcs = [getattr(mod, f"f{i}") for i in range(N)]
xs = [getattr(mod, f"x{i}") for i in range(N)]

write_module(temp_module_dir, modname, code2)
superreload(mod, old_objects) # 6.52s -> 3.24s (102% faster)

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
import importlib

Function to test (superreload and dependencies)

import io
import os
import shutil
import sys
import tempfile
import traceback
import types
import weakref

imports

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

OldObjectsMapping = dict[
tuple[str, str], list[weakref.ref]
]
from marimo._runtime.reload.autoreload import superreload

Helper for creating a temporary test module

def create_temp_module(module_name, code):
temp_dir = tempfile.mkdtemp()
module_path = os.path.join(temp_dir, f"{module_name}.py")
with open(module_path, "w", encoding="utf-8") as f:
f.write(code)
sys.path.insert(0, temp_dir)
return temp_dir, module_path

def remove_temp_module(temp_dir, module_name):
sys.path = [p for p in sys.path if p != temp_dir]
if module_name in sys.modules:
del sys.modules[module_name]
shutil.rmtree(temp_dir)

------------------ UNIT TESTS ------------------

Basic Test Cases

def test_superreload_function_reloads_module_and_updates_function():
"""
Test that superreload reloads the module and updates function code.
"""
code1 = "def foo():\n return 1\n"
code2 = "def foo():\n return 2\n"
temp_dir, module_path = create_temp_module("testmod1", code1)
try:
import testmod1
old_objects = {}
# Change code
with open(module_path, "w", encoding="utf-8") as f:
f.write(code2)
codeflash_output = superreload(testmod1, old_objects); mod = codeflash_output
finally:
remove_temp_module(temp_dir, "testmod1")

def test_superreload_class_dict_update():
"""
Test that superreload updates class attributes.
"""
code1 = "class A:\n x = 1\n"
code2 = "class A:\n x = 42\n"
temp_dir, module_path = create_temp_module("testmod2", code1)
try:
import testmod2
old_objects = {}
a = testmod2.A()
with open(module_path, "w", encoding="utf-8") as f:
f.write(code2)
codeflash_output = superreload(testmod2, old_objects); mod = codeflash_output
finally:
remove_temp_module(temp_dir, "testmod2")

def test_superreload_property_update():
"""
Test that superreload updates property objects.
"""
code1 = (
"class B:\n"
" @Property\n"
" def val(self):\n"
" return 1\n"
)
code2 = (
"class B:\n"
" @Property\n"
" def val(self):\n"
" return 99\n"
)
temp_dir, module_path = create_temp_module("testmod3", code1)
try:
import testmod3
old_objects = {}
b = testmod3.B()
with open(module_path, "w", encoding="utf-8") as f:
f.write(code2)
codeflash_output = superreload(testmod3, old_objects); mod = codeflash_output
finally:
remove_temp_module(temp_dir, "testmod3")

def test_superreload_returns_module():
"""
Test that superreload returns the module object.
"""
code = "x = 5\n"
temp_dir, module_path = create_temp_module("testmod4", code)
try:
import testmod4
old_objects = {}
codeflash_output = superreload(testmod4, old_objects); mod = codeflash_output
finally:
remove_temp_module(temp_dir, "testmod4")

Edge Test Cases

def test_superreload_with_none_old_objects():
"""
Test that superreload works if old_objects is None.
"""
code = "y = 42\n"
temp_dir, module_path = create_temp_module("testmod5", code)
try:
import testmod5
codeflash_output = superreload(testmod5, None); mod = codeflash_output
finally:
remove_temp_module(temp_dir, "testmod5")

def test_superreload_handles_module_reload_exception():
"""
Test that superreload raises if module reload fails (e.g. syntax error).
"""
code1 = "z = 1\n"
code2 = "def oops\n pass\n"
temp_dir, module_path = create_temp_module("testmod6", code1)
try:
import testmod6
old_objects = {}
with open(module_path, "w", encoding="utf-8") as f:
f.write(code2)
with pytest.raises(SyntaxError):
superreload(testmod6, old_objects)
finally:
remove_temp_module(temp_dir, "testmod6")

def test_superreload_removes_deleted_class_attributes():
"""
Test that deleted class attributes are removed after reload.
"""
code1 = "class C:\n x = 1\n y = 2\n"
code2 = "class C:\n x = 10\n"
temp_dir, module_path = create_temp_module("testmod7", code1)
try:
import testmod7
old_objects = {}
c = testmod7.C()
with open(module_path, "w", encoding="utf-8") as f:
f.write(code2)
codeflash_output = superreload(testmod7, old_objects); mod = codeflash_output
finally:
remove_temp_module(temp_dir, "testmod7")

def test_superreload_handles_non_module_objects():
"""
Test that append_obj returns False for objects not belonging to module.
"""
code = "import math\n"
temp_dir, module_path = create_temp_module("testmod8", code)
try:
import testmod8
d = {}
finally:
remove_temp_module(temp_dir, "testmod8")

def test_superreload_handles_unweakrefable_objects():
"""
Test that append_obj does not raise for objects that cannot be weakref'd.
"""
code = "x = 123\n"
temp_dir, module_path = create_temp_module("testmod9", code)
try:
import testmod9
d = {}
finally:
remove_temp_module(temp_dir, "testmod9")

def test_superreload_module_with_no_loader():
"""
Test that superreload does not raise if loader is missing.
"""
code = "a = 1\n"
temp_dir, module_path = create_temp_module("testmod10", code)
try:
import testmod10
del testmod10.loader
codeflash_output = superreload(testmod10, None); mod = codeflash_output
finally:
remove_temp_module(temp_dir, "testmod10")

Large Scale Test Cases

def test_superreload_large_module_many_functions_and_classes():
"""
Test superreload on a module with many functions and classes.
"""
N = 200
code1 = "\n".join(
[f"def f{i}(): return {i}" for i in range(N)]
+ [f"class C{i}: x = {i}" for i in range(N)]
)
code2 = "\n".join(
[f"def f{i}(): return {i+1}" for i in range(N)]
+ [f"class C{i}: x = {i+1}" for i in range(N)]
)
temp_dir, module_path = create_temp_module("testmod11", code1)
try:
import testmod11
old_objects = {}
# Check initial values
for i in range(N):
pass
# Change code
with open(module_path, "w", encoding="utf-8") as f:
f.write(code2)
codeflash_output = superreload(testmod11, old_objects); mod = codeflash_output
for i in range(N):
pass
finally:
remove_temp_module(temp_dir, "testmod11")

def test_superreload_large_module_with_properties():
"""
Test superreload on a module with many properties.
"""
N = 100
code1 = "\n".join(
[f"class D{i}:\n @Property\n def val(self):\n return {i}" for i in range(N)]
)
code2 = "\n".join(
[f"class D{i}:\n @Property\n def val(self):\n return {i+10}" for i in range(N)]
)
temp_dir, module_path = create_temp_module("testmod12", code1)
try:
import testmod12
old_objects = {}
for i in range(N):
pass
with open(module_path, "w", encoding="utf-8") as f:
f.write(code2)
codeflash_output = superreload(testmod12, old_objects); mod = codeflash_output
for i in range(N):
pass
finally:
remove_temp_module(temp_dir, "testmod12")

def test_superreload_large_module_deletes_and_adds_names():
"""
Test superreload on a module that deletes and adds many names.
"""
N = 150
code1 = "\n".join(
[f"x{i} = {i}" for i in range(N)]
)
code2 = "\n".join(
[f"y{i} = {i*2}" for i in range(N)]
)
temp_dir, module_path = create_temp_module("testmod13", code1)
try:
import testmod13
old_objects = {}
for i in range(N):
pass
with open(module_path, "w", encoding="utf-8") as f:
f.write(code2)
codeflash_output = superreload(testmod13, old_objects); mod = codeflash_output
for i in range(N):
pass
finally:
remove_temp_module(temp_dir, "testmod13")

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-mhvnfh1n and push.