⚡️ Speed up method ScopedVisitor.visit_ImportFrom by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for ScopedVisitor.visit_ImportFrom in marimo/_ast/visitor.py

⏱️ Runtime : 6.36 milliseconds → 5.75 milliseconds (best of 64 runs)

📝 Explanation and details

The optimization achieves a 10% speedup by reducing redundant operations and improving memory locality in AST processing code, specifically targeting the visit_ImportFrom method which is called frequently during import statement parsing.

Key optimizations applied:

1. Eliminated redundant string splitting in _get_alias_name: Instead of always calling node.name.split(".")[0], the code now checks for dots first and uses string slicing (name[:name.index('.')]) only when needed. This avoids creating temporary lists for simple names.

2. Cached attribute lookups: The optimization extracts self.block_stack[-1].global_names into a local variable in _define() to avoid repeated attribute traversals, and similarly caches method references (_get_alias_name, _define) and class constructors (ImportData_, VariableData_) in visit_ImportFrom.

3. Reduced string concatenation overhead: Pre-computes module_dot = module + "." once instead of concatenating inside the loop for each import.

4. Streamlined error handling: Simplified the line number extraction logic for ImportStarError by using getattr() with defaults instead of multiple hasattr() checks.

5. Loop restructuring: Split the main loop into two phases - first collecting import data, then processing definitions - which improves data locality and reduces function call overhead during the hot path.

The optimizations are particularly effective for large-scale import scenarios as shown in the test results:

• test_large_many_imports (1000 imports): 9.98% faster
• test_large_many_imports_with_as (500 imports): 10.0% faster
• test_large_mangled_locals (200 imports): 9.20% faster

For smaller import cases, the overhead of the additional setup slightly reduces performance (1-17% slower), but the dramatic improvements on large imports make this optimization worthwhile since import processing is typically dominated by large module imports in real codebases.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import ast
import uuid
from typing import Any

# imports
import pytest
from marimo._ast.visitor import ScopedVisitor


# Exceptions and helpers from marimo._ast.errors and marimo._ast.variables
class ImportStarError(Exception):
    pass

# Minimal VariableData, ImportData, Block, RefData, ObscuredScope for tests
class ImportData:
    def __init__(self, module, definition, imported_symbol, import_level):
        self.module = module
        self.definition = definition
        self.imported_symbol = imported_symbol
        self.import_level = import_level

    def __eq__(self, other):
        return (
            isinstance(other, ImportData)
            and self.module == other.module
            and self.definition == other.definition
            and self.imported_symbol == other.imported_symbol
            and self.import_level == other.import_level
        )

# Helper to extract variable_data from visitor after visit_ImportFrom
def get_variable_data(visitor: ScopedVisitor) -> dict:
    # Only look at top-level block
    return visitor.block_stack[0].variable_data

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

def test_basic_single_import():
    # Test importing a single symbol from a module
    node = ast.ImportFrom(module="math", names=[ast.alias(name="sqrt", asname=None)], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 6.57μs -> 6.84μs (3.88% slower)
    data = get_variable_data(visitor)
    vd = data["sqrt"][0]

def test_basic_import_with_as():
    # Test importing with alias
    node = ast.ImportFrom(module="math", names=[ast.alias(name="sqrt", asname="mysqrt")], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 6.71μs -> 7.06μs (4.93% slower)
    data = get_variable_data(visitor)
    vd = data["mysqrt"][0]

def test_basic_multiple_imports():
    # Test importing multiple symbols
    node = ast.ImportFrom(
        module="math",
        names=[ast.alias(name="sin", asname=None), ast.alias(name="cos", asname=None)],
        level=0,
    )
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 8.23μs -> 8.34μs (1.31% slower)
    data = get_variable_data(visitor)

def test_basic_relative_import():
    # Test relative import (level > 0)
    node = ast.ImportFrom(module="foo", names=[ast.alias(name="bar", asname=None)], level=2)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 5.54μs -> 5.87μs (5.59% slower)
    data = get_variable_data(visitor)

def test_basic_no_module():
    # Test ImportFrom with module=None (should use "")
    node = ast.ImportFrom(module=None, names=[ast.alias(name="baz", asname=None)], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 5.50μs -> 5.79μs (4.91% slower)
    data = get_variable_data(visitor)

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

def test_import_star_raises():
    # Test that import * raises ImportStarError
    node = ast.ImportFrom(module="math", names=[ast.alias(name="*", asname=None)], level=0)
    visitor = ScopedVisitor()
    with pytest.raises(ImportStarError) as excinfo:
        visitor.visit_ImportFrom(node)

def test_import_as_local_mangles():
    # Test that asname that is local gets mangled
    node = ast.ImportFrom(module="pkg", names=[ast.alias(name="foo", asname="_bar")], level=0)
    visitor = ScopedVisitor(mangle_prefix="abc")
    visitor.visit_ImportFrom(node) # 9.46μs -> 10.5μs (9.65% slower)
    data = get_variable_data(visitor)

def test_import_as_double_underscore_local():
    # Test that asname "__" gets mangled
    node = ast.ImportFrom(module="pkg", names=[ast.alias(name="foo", asname="__")], level=0)
    visitor = ScopedVisitor(mangle_prefix="xyz")
    visitor.visit_ImportFrom(node) # 7.12μs -> 7.22μs (1.43% slower)
    data = get_variable_data(visitor)

def test_import_as_not_local():
    # Test that asname not local does not get mangled
    node = ast.ImportFrom(module="pkg", names=[ast.alias(name="foo", asname="bar")], level=0)
    visitor = ScopedVisitor(mangle_prefix="zzz")
    visitor.visit_ImportFrom(node) # 6.35μs -> 6.78μs (6.32% slower)
    data = get_variable_data(visitor)

def test_import_dot_name():
    # Test that dotted names only define the first part
    node = ast.ImportFrom(module="pkg", names=[ast.alias(name="foo.bar.baz", asname=None)], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 5.95μs -> 7.16μs (16.9% slower)
    data = get_variable_data(visitor)

def test_import_with_lineno():
    # Test that ImportStarError includes correct line number
    node = ast.ImportFrom(module="pkg", names=[ast.alias(name="*", asname=None)], level=0)
    node.lineno = 42
    visitor = ScopedVisitor()
    with pytest.raises(ImportStarError) as excinfo:
        visitor.visit_ImportFrom(node)

def test_import_star_no_lineno():
    # Test ImportStarError with no lineno attribute
    node = ast.ImportFrom(module="pkg", names=[ast.alias(name="*", asname=None)], level=0)
    visitor = ScopedVisitor()
    with pytest.raises(ImportStarError) as excinfo:
        visitor.visit_ImportFrom(node)

def test_ignore_local_flag():
    # Test ignore_local disables mangling
    node = ast.ImportFrom(module="pkg", names=[ast.alias(name="foo", asname="_bar")], level=0)
    visitor = ScopedVisitor(mangle_prefix="abc", ignore_local=True)
    visitor.visit_ImportFrom(node) # 9.78μs -> 9.86μs (0.852% slower)
    data = get_variable_data(visitor)

# --------------------------
# Large Scale Test Cases
# --------------------------

def test_large_many_imports():
    # Test importing many symbols at once (1000)
    names = [ast.alias(name=f"sym{i}", asname=None) for i in range(1000)]
    node = ast.ImportFrom(module="bigmod", names=names, level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 1.45ms -> 1.31ms (9.98% faster)
    data = get_variable_data(visitor)
    for i in range(1000):
        key = f"sym{i}"

def test_large_many_imports_with_as():
    # Test importing many symbols with asname (500)
    names = [ast.alias(name=f"sym{i}", asname=f"alias{i}") for i in range(500)]
    node = ast.ImportFrom(module="bigmod", names=names, level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 797μs -> 724μs (10.0% faster)
    data = get_variable_data(visitor)
    for i in range(500):
        key = f"alias{i}"

def test_large_mangled_locals():
    # Test importing many local asnames, all should be mangled
    names = [ast.alias(name=f"sym{i}", asname=f"_local{i}") for i in range(200)]
    node = ast.ImportFrom(module="bigmod", names=names, level=0)
    visitor = ScopedVisitor(mangle_prefix="mng")
    visitor.visit_ImportFrom(node) # 352μs -> 323μs (9.20% faster)
    data = get_variable_data(visitor)
    for i in range(200):
        key = f"_mng_local{i}"

def test_large_mix_local_and_nonlocal():
    # Mix local and non-local asnames
    names = []
    for i in range(250):
        if i % 2 == 0:
            names.append(ast.alias(name=f"sym{i}", asname=f"_loc{i}"))
        else:
            names.append(ast.alias(name=f"sym{i}", asname=f"alias{i}"))
    node = ast.ImportFrom(module="bigmod", names=names, level=0)
    visitor = ScopedVisitor(mangle_prefix="mix")
    visitor.visit_ImportFrom(node) # 424μs -> 380μs (11.6% faster)
    data = get_variable_data(visitor)
    for i in range(250):
        if i % 2 == 0:
            key = f"_mix_loc{i}"
        else:
            key = f"alias{i}"
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import ast

# imports
import pytest
from marimo._ast.visitor import ScopedVisitor

# Function and dependencies to test (from marimo/_ast/visitor.py and marimo/_ast/variables.py)

class ImportStarError(Exception):
    pass

class ImportData:
    def __init__(self, module, definition, imported_symbol, import_level):
        self.module = module
        self.definition = definition
        self.imported_symbol = imported_symbol
        self.import_level = import_level

    def __eq__(self, other):
        return (
            isinstance(other, ImportData)
            and self.module == other.module
            and self.definition == other.definition
            and self.imported_symbol == other.imported_symbol
            and self.import_level == other.import_level
        )

# Helper function to extract defined names and variable data from visitor after import
def get_import_defs(visitor):
    block = visitor.block_stack[-1]
    return block.defs, block.variable_data

# ------------------- UNIT TESTS -------------------

# Basic Test Cases

def test_basic_single_import_from():
    # from math import sqrt
    node = ast.ImportFrom(module="math", names=[ast.alias(name="sqrt", asname=None)], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 5.81μs -> 7.06μs (17.7% slower)
    defs, vardata = get_import_defs(visitor)

def test_basic_multiple_import_from():
    # from os.path import join, dirname
    node = ast.ImportFrom(module="os.path", names=[
        ast.alias(name="join", asname=None),
        ast.alias(name="dirname", asname=None)
    ], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 8.43μs -> 8.41μs (0.226% faster)
    defs, vardata = get_import_defs(visitor)

def test_basic_import_with_as():
    # from sys import version as v
    node = ast.ImportFrom(module="sys", names=[ast.alias(name="version", asname="v")], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 6.44μs -> 6.74μs (4.43% slower)
    defs, vardata = get_import_defs(visitor)

def test_basic_import_level():
    # from ..foo import bar
    node = ast.ImportFrom(module="foo", names=[ast.alias(name="bar", asname=None)], level=2)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 5.45μs -> 5.76μs (5.40% slower)
    defs, vardata = get_import_defs(visitor)

# Edge Test Cases

def test_import_from_with_none_module():
    # from . import something
    node = ast.ImportFrom(module=None, names=[ast.alias(name="something", asname=None)], level=1)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 5.28μs -> 5.83μs (9.38% slower)
    defs, vardata = get_import_defs(visitor)

def test_import_from_with_dotted_name():
    # from pkg import sub.mod
    node = ast.ImportFrom(module="pkg", names=[ast.alias(name="sub.mod", asname=None)], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 5.86μs -> 7.08μs (17.2% slower)
    defs, vardata = get_import_defs(visitor)

def test_import_from_with_as_and_local_mangling():
    # from foo import bar as _baz (should mangle _baz)
    node = ast.ImportFrom(module="foo", names=[ast.alias(name="bar", asname="_baz")], level=0)
    visitor = ScopedVisitor(mangle_prefix="XyZ_")
    visitor.visit_ImportFrom(node) # 7.12μs -> 7.17μs (0.697% slower)
    defs, vardata = get_import_defs(visitor)

def test_import_from_with_as_and_double_underscore_local():
    # from foo import bar as __
    node = ast.ImportFrom(module="foo", names=[ast.alias(name="bar", asname="__")], level=0)
    visitor = ScopedVisitor(mangle_prefix="Cell_")
    visitor.visit_ImportFrom(node) # 6.31μs -> 6.65μs (5.19% slower)
    defs, vardata = get_import_defs(visitor)

def test_import_from_with_as_and_double_underscore_not_local():
    # from foo import bar as "__baz"
    node = ast.ImportFrom(module="foo", names=[ast.alias(name="bar", asname="__baz")], level=0)
    visitor = ScopedVisitor(mangle_prefix="Cell_")
    visitor.visit_ImportFrom(node) # 6.39μs -> 6.53μs (2.14% slower)
    defs, vardata = get_import_defs(visitor)

def test_import_star_raises():
    # from foo import *
    node = ast.ImportFrom(module="foo", names=[ast.alias(name="*", asname=None)], level=0)
    visitor = ScopedVisitor()
    with pytest.raises(ImportStarError) as excinfo:
        visitor.visit_ImportFrom(node)

def test_import_from_with_empty_names():
    # from foo import (nothing)
    node = ast.ImportFrom(module="foo", names=[], level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 846ns -> 1.70μs (50.2% slower)
    defs, vardata = get_import_defs(visitor)

def test_import_from_with_global_names():
    # from foo import bar, baz, qux (baz is global)
    node = ast.ImportFrom(module="foo", names=[
        ast.alias(name="bar", asname=None),
        ast.alias(name="baz", asname=None),
        ast.alias(name="qux", asname=None)
    ], level=0)
    visitor = ScopedVisitor()
    # Mark 'baz' as global
    visitor.block_stack[-1].global_names.add("baz")
    visitor.visit_ImportFrom(node) # 12.2μs -> 12.0μs (1.22% faster)
    # 'baz' should be defined in block 0, others in -1
    defs0 = visitor.block_stack[0].defs
    defs1 = visitor.block_stack[-1].defs

# Large Scale Test Cases

def test_large_scale_many_imports():
    # from foo import a1, a2, ..., a999
    names = [ast.alias(name=f"a{i}", asname=None) for i in range(1, 1000)]
    node = ast.ImportFrom(module="foo", names=names, level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 1.44ms -> 1.30ms (10.5% faster)
    defs, vardata = get_import_defs(visitor)
    # All a1..a999 should be defined
    for i in range(1, 1000):
        name = f"a{i}"

def test_large_scale_many_imports_with_as_and_local():
    # from foo import a1 as _b1, a2 as _b2, ..., a999 as _b999
    names = [ast.alias(name=f"a{i}", asname=f"_b{i}") for i in range(1, 1000)]
    node = ast.ImportFrom(module="foo", names=names, level=0)
    visitor = ScopedVisitor(mangle_prefix="TestPrefix_")
    visitor.visit_ImportFrom(node) # 1.73ms -> 1.53ms (13.5% faster)
    defs, vardata = get_import_defs(visitor)
    for i in range(1, 1000):
        mangled = f"_TestPrefix__b{i}"

def test_large_scale_import_star_raises_first():
    # from foo import *, a1, a2, ..., a999
    names = [ast.alias(name="*", asname=None)] + [ast.alias(name=f"a{i}", asname=None) for i in range(1, 10)]
    node = ast.ImportFrom(module="foo", names=names, level=0)
    visitor = ScopedVisitor()
    with pytest.raises(ImportStarError):
        visitor.visit_ImportFrom(node)

def test_large_scale_import_star_raises_last():
    # from foo import a1, a2, ..., a10, *
    names = [ast.alias(name=f"a{i}", asname=None) for i in range(1, 10)] + [ast.alias(name="*", asname=None)]
    node = ast.ImportFrom(module="foo", names=names, level=0)
    visitor = ScopedVisitor()
    with pytest.raises(ImportStarError):
        visitor.visit_ImportFrom(node)

def test_large_scale_import_from_with_mixed_as_and_plain():
    # from foo import a1, a2 as b2, a3, a4 as b4, ..., a10 as b10
    names = []
    for i in range(1, 11):
        if i % 2 == 0:
            names.append(ast.alias(name=f"a{i}", asname=f"b{i}"))
        else:
            names.append(ast.alias(name=f"a{i}", asname=None))
    node = ast.ImportFrom(module="foo", names=names, level=0)
    visitor = ScopedVisitor()
    visitor.visit_ImportFrom(node) # 25.9μs -> 26.0μs (0.422% slower)
    defs, vardata = get_import_defs(visitor)
    for i in range(1, 11):
        if i % 2 == 0:
            pass
        else:
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from ast import ImportFrom
from marimo._ast.visitor import ScopedVisitor
import pytest

def test_ScopedVisitor_visit_ImportFrom():
    with pytest.raises(AttributeError, match="'ImportFrom'\\ object\\ has\\ no\\ attribute\\ 'names'"):
        ScopedVisitor.visit_ImportFrom(ScopedVisitor(), ImportFrom())

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_o_lbxivc/tmpozcyxcqt/test_concolic_coverage.py::test_ScopedVisitor_visit_ImportFrom	1.26μs	1.36μs	-7.78%⚠️

To edit these changes git checkout codeflash/optimize-ScopedVisitor.visit_ImportFrom-mhcypelv and push.