⚡️ Speed up function _format_variables by 23%

[codeflash-ai]

📄 23% (0.23x) speedup for _format_variables in marimo/_server/ai/prompts.py

⏱️ Runtime : 496 microseconds → 404 microseconds (best of 45 runs)

📝 Explanation and details

The optimization replaces inefficient string concatenation with a list-based approach, delivering a 22% speedup. Here's what changed:

Key Optimization: String Building Strategy

• Original: Used variable_info += f"..." to repeatedly concatenate strings in a loop
• Optimized: Builds a list with lines.append() and joins once at the end with "".join(lines)

Why This is Faster:
String concatenation with += in Python creates a new string object each time, copying all previous content. With N variables, this results in O(N²) time complexity and excessive memory allocations. The list approach is O(N) - each string is added once, then efficiently joined.

Minor Enhancement:
Cached lines.append as a local variable append to avoid repeated attribute lookups during the loop, providing additional microsecond-level gains.

Performance Profile:

• Line profiler shows the optimized version reduces time spent on string operations (line with variable_info += went from 29.8% to 25.8% of total time)
• The optimization shines with larger datasets: test cases with 1000 variables show 21-29% speedup, while smaller cases show minimal gains or slight overhead due to list setup costs
• Best suited for scenarios with multiple variables to format, where the O(N²) vs O(N) difference becomes significant

The code maintains identical output format and error handling behavior.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Union

# imports
import pytest  # used for our unit tests
from marimo._server.ai.prompts import _format_variables

# function to test
# Copyright 2024 Marimo. All rights reserved.


# Minimal VariableContext definition for testing purposes
@dataclass
class VariableContext:
    name: str
    value_type: str
    preview_value: str
from marimo._server.ai.prompts import _format_variables

# unit tests

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

def test_empty_list_returns_empty_string():
    # Test with empty list
    codeflash_output = _format_variables([]) # 376ns -> 466ns (19.3% slower)
    # Test with None
    codeflash_output = _format_variables(None) # 171ns -> 202ns (15.3% slower)


def test_single_string_variable_basic():
    # Test with one string variable, non-private
    expected = (
        "\n\n## Available variables from other cells:\n"
        "- variable: `bar`"
    )
    codeflash_output = _format_variables(["bar"]) # 1.51μs -> 2.06μs (26.7% slower)




def test_private_string_variable_skipped():
    # String variable with private name should be skipped
    codeflash_output = _format_variables(["_private"]) # 1.20μs -> 1.66μs (27.7% slower)




def test_string_variable_with_empty_string():
    # Empty string variable (not private), should appear
    expected = (
        "\n\n## Available variables from other cells:\n"
        "- variable: ``"
    )
    codeflash_output = _format_variables([""]) # 1.46μs -> 2.02μs (27.5% slower)


def test_string_variable_with_private_like_name():
    # String variable with "_" but not starting with "_"
    expected = (
        "\n\n## Available variables from other cells:\n"
        "- variable: `foo_bar`"
    )
    codeflash_output = _format_variables(["foo_bar"]) # 1.52μs -> 2.13μs (28.5% slower)

def test_variablecontext_with_non_string_name():
    # VariableContext with name as non-string (should raise AttributeError)
    var = VariableContext(name=123, value_type="int", preview_value="123")
    with pytest.raises(AttributeError):
        _format_variables([var]) # 1.83μs -> 1.91μs (4.35% slower)

def test_string_variable_with_non_string_type():
    # String variable is not a string (should raise AttributeError)
    with pytest.raises(AttributeError):
        _format_variables([123]) # 1.59μs -> 1.79μs (11.4% slower)



def test_large_number_of_string_variables():
    # Test with 1000 string variables
    variables = [f"var{i}" for i in range(1000)]
    codeflash_output = _format_variables(variables); result = codeflash_output # 156μs -> 128μs (21.1% faster)
    # Should contain all variable names
    for i in range(1000):
        pass



def test_large_all_empty_strings():
    # 1000 empty string variables
    variables = [""] * 1000
    codeflash_output = _format_variables(variables); result = codeflash_output # 160μs -> 124μs (28.9% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Union

# imports
import pytest  # used for our unit tests
from marimo._server.ai.prompts import _format_variables

# function to test
# Copyright 2024 Marimo. All rights reserved.


@dataclass
class VariableContext:
    name: str
    value_type: str
    preview_value: str
from marimo._server.ai.prompts import _format_variables

# unit tests

# ------------- BASIC TEST CASES -------------

def test_empty_list_returns_empty_string():
    # Test with empty list
    codeflash_output = _format_variables([]) # 411ns -> 401ns (2.49% faster)

def test_none_returns_empty_string():
    # Test with None
    codeflash_output = _format_variables(None) # 363ns -> 371ns (2.16% slower)


def test_single_string_variable():
    # Test with a single string variable
    var = "bar"
    expected = (
        "\n\n## Available variables from other cells:\n"
        "- variable: `bar`"
    )
    codeflash_output = _format_variables([var]) # 1.55μs -> 2.04μs (24.0% slower)


def test_multiple_string_variables():
    # Test with multiple string variables
    vars = ["a", "b", "c"]
    expected = (
        "\n\n## Available variables from other cells:\n"
        "- variable: `a`"
        "- variable: `b`"
        "- variable: `c`"
    )
    codeflash_output = _format_variables(vars) # 2.13μs -> 2.66μs (20.0% slower)



def test_private_string_variable_ignored():
    # Test that string variable with private name is ignored
    var = "_private"
    codeflash_output = _format_variables([var]) # 1.18μs -> 1.69μs (30.3% slower)







def test_variablecontext_with_non_str_types():
    # VariableContext fields are not strictly strings (should be handled as str)
    var = VariableContext(name=123, value_type=456, preview_value=789)
    # The function expects .name, .value_type, .preview_value are strings
    # This will raise AttributeError if not handled, so test that it fails
    with pytest.raises(AttributeError):
        _format_variables([var]) # 1.99μs -> 2.24μs (11.2% slower)

def test_string_variable_is_not_str():
    # Variable is not VariableContext and not a string (should raise AttributeError)
    class Dummy:
        pass
    with pytest.raises(AttributeError):
        _format_variables([Dummy()]) # 1.88μs -> 1.97μs (4.76% slower)

# ------------- LARGE SCALE TEST CASES -------------

def test_large_number_of_public_string_variables():
    # Test with 1000 public string variables
    vars = [f"var{i}" for i in range(1000)]
    codeflash_output = _format_variables(vars); result = codeflash_output # 158μs -> 124μs (27.3% faster)
    # Should contain all variable names
    for i in range(1000):
        pass




#------------------------------------------------
from marimo._server.ai.prompts import _format_variables

def test__format_variables():
    _format_variables(['', '_'])

def test__format_variables_2():
    _format_variables([])

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_qdbuqf4z/tmpagq95vjv/test_concolic_coverage.py::test__format_variables	1.65μs	2.06μs	-19.9%⚠️
codeflash_concolic_qdbuqf4z/tmpagq95vjv/test_concolic_coverage.py::test__format_variables_2	425ns	385ns	10.4%✅

To edit these changes git checkout codeflash/optimize-_format_variables-mh5jt82w and push.