-
Notifications
You must be signed in to change notification settings - Fork 21
alias support and vitest imports #1266
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
⚡️ Codeflash found optimizations for this PR📄 42% (0.42x) speedup for
|
| needs_import = any(f"{global_name}(" in code or f"{global_name} (" in code for global_name in test_globals) | ||
|
|
||
| if not needs_import: | ||
| return code | ||
|
|
||
| # Determine which globals are actually used in the code | ||
| used_globals = [g for g in test_globals if f"{g}(" in code or f"{g} (" in code] | ||
| if not used_globals: | ||
| return code | ||
|
|
||
| # Build the import statement | ||
| import_statement = f"import {{ {', '.join(used_globals)} }} from 'vitest';\n" | ||
|
|
||
| # Find the first line that isn't a comment or empty | ||
| lines = code.split("\n") | ||
| insert_index = 0 | ||
| for i, line in enumerate(lines): | ||
| stripped = line.strip() | ||
| if stripped and not stripped.startswith("//") and not stripped.startswith("/*") and not stripped.startswith("*"): | ||
| # Check if this line is an import/require - insert after imports | ||
| if stripped.startswith("import ") or stripped.startswith("const ") or stripped.startswith("let "): | ||
| continue | ||
| insert_index = i | ||
| break | ||
| insert_index = i + 1 | ||
|
|
||
| # Find the last import line to insert after it | ||
| last_import_index = -1 | ||
| for i, line in enumerate(lines): | ||
| stripped = line.strip() | ||
| if stripped.startswith("import ") and "from " in stripped: | ||
| last_import_index = i | ||
|
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
⚡️Codeflash found 23% (0.23x) speedup for ensure_vitest_imports in codeflash/languages/javascript/module_system.py
⏱️ Runtime : 828 microseconds → 673 microseconds (best of 250 runs)
📝 Explanation and details
The optimized code achieves a 23% runtime improvement by eliminating redundant scanning of the code string and consolidating loop iterations over the code lines.
Key optimizations:
-
Single-pass global detection: The original code scanned for test globals twice—once with
any()to check if imports are needed, then again with a list comprehension to collect used globals. The optimized version computesused_globalsin a single pass and derivesneeds_importfrom whether that list is empty usingbool(used_globals). This eliminates one complete scan through the code string for all 9 test globals. -
Merged loop for insertion logic: The original code had two separate loops over
lines:- First loop: Find
insert_index(where to insert if no imports exist) - Second loop: Find
last_import_index(the last existing import line)
The optimized version merges these into a single loop, checking for both conditions simultaneously. This cuts the number of line iterations in half, reducing overhead from
enumerate(),strip()calls, and string comparisons. - First loop: Find
Performance characteristics:
The test results show the optimization is most effective for:
- Large files (27-41% faster): The benefit compounds with file size since we avoid redundant string scans and line iterations
- Multiple globals usage (30-48% faster): Cases using
test,expect,it,afterAll, etc. benefit significantly from the single-pass detection - Files with many imports (10-18% faster): Merged loop reduces overhead when scanning through existing import statements
The optimization maintains identical behavior—all test cases pass with the same output while running faster across all scenarios, from simple cases (~6-35% faster) to complex nested structures (~41% faster).
✅ Correctness verification report:
| Test | Status |
|---|---|
| ⚙️ Existing Unit Tests | 🔘 None Found |
| 🌀 Generated Regression Tests | ✅ 49 Passed |
| ⏪ Replay Tests | 🔘 None Found |
| 🔎 Concolic Coverage Tests | 🔘 None Found |
| 📊 Tests Coverage | 96.9% |
🌀 Click to see Generated Regression Tests
import logging
# imports
import pytest # used for our unit tests
from codeflash.languages.javascript.module_system import ensure_vitest_imports
def test_non_vitest_framework_returns_same():
# If the framework is not 'vitest', the code must be returned unchanged.
original = "describe('a', () => {});"
codeflash_output = ensure_vitest_imports(original, test_framework="jest"); result = codeflash_output # 861ns -> 842ns (2.26% faster)
def test_no_globals_no_insertion():
# When the code doesn't use any vitest globals, nothing should be added.
original = "// file with comments only\nconst x = 1;\nfunction foo() { return x; }\n"
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 4.98μs -> 4.59μs (8.50% faster)
def test_already_imported_single_quotes():
# If the code already imports from 'vitest' (single quotes), return unchanged.
original = "import { describe } from 'vitest';\ndescribe('a', () => {});"
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 1.04μs -> 981ns (6.22% faster)
def test_already_imported_double_quotes():
# If the code already imports from "vitest" (double quotes), return unchanged.
original = 'import { test } from "vitest";\ntest("works", () => {});'
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 1.06μs -> 1.05μs (0.951% faster)
def test_basic_add_import_at_top_for_describe_and_test():
# When using vitest and code contains globals, the function should add the import.
original = "describe('suite', () => {\n test('case', () => {});\n});"
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 10.4μs -> 8.45μs (22.9% faster)
def test_insert_after_last_import_line():
# If the code has existing import ... from ... lines, the vitest import
# should be inserted after the last such import.
original = (
"import foo from 'foo';\n"
"import bar from 'bar';\n"
"// some comment\n"
"describe('a', () => {});\n"
)
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 11.7μs -> 9.95μs (17.3% faster)
lines = result.splitlines()
def test_insert_after_const_or_let_lines():
# Lines beginning with const/let should be skipped when determining insertion
# point, so the vitest import should be inserted after them.
original = "const x = require('x');\ndescribe('a', () => {});"
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 9.92μs -> 8.32μs (19.1% faster)
lines = result.splitlines()
def test_detects_usage_in_string_literals_and_inserts():
# The function currently checks substrings, so if a global name appears inside
# a string literal like "describe(", it will still trigger an import.
original = "const s = \"this mentions describe(\";\nconsole.log(s);"
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 9.91μs -> 8.28μs (19.7% faster)
def test_handles_space_before_parenthesis():
# The detection should handle both 'it(' and 'it (' (with a space).
original = "it ('should pass', () => {});"
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 9.73μs -> 7.39μs (31.6% faster)
def test_used_globals_order_is_preserved():
# When multiple globals are used, their order in the import should follow
# the original test_globals list defined in the function.
# Here we use 'vi' and 'expect' so import should be "expect, vi"
original = "vi(functionStub);\nexpect(1).toBe(1);\n"
# Note: order in the code is vi first, expect second, but used_globals must follow
# the fixed ordering inside the function (expect before vi).
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 10.4μs -> 7.66μs (35.3% faster)
# Find the import line
import_lines = [l for l in result.splitlines() if l.startswith("import {")]
import_line = import_lines[0]
def test_large_scale_insertion_efficiency_and_correctness():
# Large-scale-ish test: create a file with several hundred lines (under 1000)
# and a few scattered usages of vitest globals. Ensure the import is added
# exactly once and lists only the used globals (in the canonical order).
fragments = []
# Add 250 noop lines
for i in range(250):
fragments.append(f"// filler line {i}")
# Add some import statements to ensure insertion after last import
fragments.append("import alpha from 'alpha';")
fragments.append("import beta from 'beta';")
# Add more filler
for i in range(250, 500):
fragments.append(f"const val{i} = {i};")
# Use a few vitest globals scattered through the file
fragments.append("beforeEach(() => {});")
fragments.append("someOtherLine();")
fragments.append("afterAll(() => {});")
fragments.append("console.log('done');")
# Join into a big code string
big_code = "\n".join(fragments)
codeflash_output = ensure_vitest_imports(big_code, test_framework="vitest"); result = codeflash_output # 346μs -> 271μs (27.4% faster)
# Ensure the import was inserted after the last explicit import ('import beta from ...')
lines = result.splitlines()
# Find index of "import beta from 'beta';"
beta_idx = next(i for i, l in enumerate(lines) if l.strip() == "import beta from 'beta';")
def test_no_duplicate_insertion_on_multiple_global_occurrences():
# If globals appear multiple times, the import should still list each used global once.
original = (
"describe('one', () => {});\n"
"describe('two', () => {});\n"
"test('a', () => {});\n"
"test('b', () => {});\n"
)
codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 10.9μs -> 8.73μs (25.0% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.import pytest
from codeflash.languages.javascript.module_system import ensure_vitest_imports
def test_non_vitest_framework_returns_code_unchanged():
"""Test that code is returned unchanged when test_framework is not 'vitest'."""
code = "describe('test', () => { test('works', () => { expect(1).toBe(1); }); });"
codeflash_output = ensure_vitest_imports(code, "jest"); result = codeflash_output # 511ns -> 521ns (1.92% slower)
def test_jest_framework_ignores_vitest_needs():
"""Test that jest framework is not modified even if vitest imports are missing."""
code = "describe('suite', () => { test('case', () => {}); });"
codeflash_output = ensure_vitest_imports(code, "jest"); result = codeflash_output # 470ns -> 441ns (6.58% faster)
def test_mocha_framework_returns_unchanged():
"""Test that mocha framework code is returned unchanged."""
code = "describe('mocha test', () => { it('works', () => {}); });"
codeflash_output = ensure_vitest_imports(code, "mocha"); result = codeflash_output # 441ns -> 421ns (4.75% faster)
def test_vitest_with_existing_import_single_quotes():
"""Test that existing vitest imports with single quotes are recognized."""
code = "import { describe, test } from 'vitest';\ndescribe('test', () => {});"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 731ns -> 771ns (5.19% slower)
def test_vitest_with_existing_import_double_quotes():
"""Test that existing vitest imports with double quotes are recognized."""
code = 'import { describe, test } from "vitest";\ndescribe("test", () => {});'
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 822ns -> 872ns (5.73% slower)
def test_vitest_adds_import_for_describe():
"""Test that vitest import is added when describe is used."""
code = "describe('test', () => {});"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.30μs -> 7.49μs (24.1% faster)
def test_vitest_adds_import_for_test():
"""Test that vitest import is added when test is used."""
code = "test('my test', () => { expect(true).toBe(true); });"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.39μs -> 7.17μs (30.9% faster)
def test_vitest_adds_import_for_expect():
"""Test that vitest import is added when expect is used."""
code = "expect(5).toBe(5);"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.23μs -> 6.84μs (34.8% faster)
def test_vitest_adds_import_for_it():
"""Test that vitest import is added when it is used."""
code = "it('should work', () => { expect(true).toBe(true); });"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.48μs -> 7.30μs (29.8% faster)
def test_vitest_no_import_if_no_globals_used():
"""Test that no import is added if test globals are not used."""
code = "const x = 5; console.log(x);"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.82μs -> 3.56μs (7.31% faster)
def test_vitest_handles_multiple_globals():
"""Test that multiple vitest globals are included in a single import."""
code = "describe('suite', () => { test('case', () => { expect(1).toBe(1); }); });"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.26μs -> 7.61μs (21.6% faster)
def test_vitest_with_space_before_parenthesis():
"""Test that globals with space before parenthesis are recognized."""
code = "describe ('test', () => {});"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 8.74μs -> 7.03μs (24.2% faster)
def test_vitest_with_multiple_spaces_before_parenthesis():
"""Test that globals with multiple spaces before parenthesis are recognized."""
code = "test ('my test', () => {});"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.90μs -> 3.66μs (6.56% faster)
def test_vitest_ignores_globals_in_comments():
"""Test that globals mentioned in comments don't trigger imports."""
code = "// describe('test', () => {});\nconst x = 5;"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.40μs -> 7.83μs (20.0% faster)
def test_vitest_ignores_globals_in_block_comments():
"""Test that globals in block comments are ignored."""
code = "/* describe('test', () => {}); */\nconst x = 5;"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 8.97μs -> 7.41μs (20.9% faster)
def test_vitest_with_leading_comments():
"""Test that import is placed correctly after leading comments."""
code = "// This is a test file\n// Another comment\ndescribe('test', () => {});"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.88μs -> 8.27μs (19.5% faster)
lines = result.split("\n")
def test_vitest_with_existing_imports_before_code():
"""Test that vitest import is inserted after existing imports."""
code = "import axios from 'axios';\ndescribe('test', () => {});"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.74μs -> 8.20μs (18.8% faster)
lines = result.split("\n")
# Vitest import should come after
vitest_line_index = next(i for i, line in enumerate(lines) if "vitest" in line)
def test_vitest_with_beforeeach():
"""Test that beforeEach is recognized and imported."""
code = "beforeEach(() => { setup(); });"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.77μs -> 6.91μs (41.3% faster)
def test_vitest_with_aftereach():
"""Test that afterEach is recognized and imported."""
code = "afterEach(() => { cleanup(); });"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.70μs -> 7.01μs (38.3% faster)
def test_vitest_with_beforeall():
"""Test that beforeAll is recognized and imported."""
code = "beforeAll(() => { init(); });"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 8.84μs -> 6.97μs (26.7% faster)
def test_vitest_with_afterall():
"""Test that afterAll is recognized and imported."""
code = "afterAll(() => { destroy(); });"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.91μs -> 6.69μs (48.1% faster)
def test_vitest_with_vi_mock():
"""Test that vi mock utility is recognized and imported."""
code = "vi.mock('./module');"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.90μs -> 3.52μs (10.8% faster)
def test_vitest_empty_code():
"""Test handling of empty code string."""
code = ""
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.01μs -> 2.69μs (12.0% faster)
def test_vitest_only_whitespace():
"""Test handling of code with only whitespace."""
code = " \n \n "
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.48μs -> 3.17μs (9.79% faster)
def test_vitest_global_not_as_function_call():
"""Test that globals not followed by parenthesis are not imported."""
code = "const test = 5; const describe = 'string';"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.90μs -> 3.63μs (7.44% faster)
def test_vitest_case_sensitive_globals():
"""Test that global matching is case-sensitive."""
code = "Describe('test', () => {}); Test('case', () => {});"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 4.27μs -> 4.08μs (4.66% faster)
def test_vitest_mixed_existing_and_needed_imports():
"""Test handling when some imports exist and some are needed."""
code = "import { describe } from 'vitest';\ntest('case', () => {});"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 701ns -> 731ns (4.10% slower)
def test_vitest_with_line_continuation():
"""Test code with line continuations and multiple statements."""
code = "describe(\n 'test',\n () => {\n test('case', () => {});\n }\n);"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 10.9μs -> 8.40μs (30.1% faster)
def test_vitest_preserves_original_content():
"""Test that the original code content is preserved after adding imports."""
code = "describe('test', () => { expect(true).toBe(true); });"
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.03μs -> 7.48μs (20.6% faster)
def test_vitest_large_code_file_with_many_tests():
"""Test handling of a large file with many test cases."""
# Generate a code string with 50 test cases
code_parts = []
for i in range(50):
code_parts.append(f"test('test case {i}', () => {{ expect({i}).toBe({i}); }});")
code = "\n".join(code_parts)
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 40.9μs -> 31.9μs (28.2% faster)
def test_vitest_large_file_with_multiple_import_types():
"""Test handling of a large file with many existing imports."""
# Create code with many imports before test code
imports = []
for i in range(50):
imports.append(f"import module{i} from 'module{i}';")
imports.append("describe('suite', () => { test('case', () => {}); });")
code = "\n".join(imports)
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 45.9μs -> 41.7μs (10.1% faster)
# All original imports should be preserved
for i in range(50):
pass
def test_vitest_large_code_with_mixed_globals():
"""Test handling of code that uses all supported vitest globals."""
globals_used = ["describe", "test", "it", "expect", "vi", "beforeEach", "afterEach", "beforeAll", "afterAll"]
code_parts = []
for global_name in globals_used:
code_parts.append(f"{global_name}(() => {{}});")
code = "\n".join(code_parts)
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 10.1μs -> 7.61μs (32.1% faster)
# All globals should be in the import
for global_name in globals_used:
pass
def test_vitest_large_file_no_modifications_needed():
"""Test performance with large file that doesn't need modifications."""
# Create a large file with no vitest globals
code_parts = []
for i in range(100):
code_parts.append(f"const var{i} = {i};")
code = "\n".join(code_parts)
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 19.4μs -> 19.1μs (1.79% faster)
def test_vitest_complex_nested_code_structure():
"""Test handling of deeply nested code structures with multiple globals."""
code = """
describe('outer suite', () => {
beforeAll(() => {});
afterAll(() => {});
describe('inner suite', () => {
beforeEach(() => {});
afterEach(() => {});
test('test 1', () => {
expect(1).toBe(1);
vi.mock('./mod');
});
it('test 2', () => {
expect(2).toBe(2);
});
});
});
"""
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 15.7μs -> 11.1μs (41.0% faster)
def test_vitest_code_with_long_lines():
"""Test handling of code with very long lines."""
long_line = "describe('test', () => { " + "test('case', () => { expect(x).toBe(y); }); " * 50 + "});"
code = long_line
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 22.4μs -> 20.7μs (8.08% faster)
def test_vitest_performance_with_many_comment_lines():
"""Test that performance is acceptable with many comment lines."""
comments = ["// This is a comment"] * 100
code_lines = comments + ["describe('test', () => { test('case', () => {}); });"]
code = "\n".join(code_lines)
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 52.8μs -> 46.6μs (13.3% faster)
def test_vitest_repeated_globals_not_duplicated():
"""Test that repeated use of same global doesn't create duplicates in import."""
code = """
describe('suite', () => {});
describe('another suite', () => {});
test('case1', () => {});
test('case2', () => {});
test('case3', () => {});
"""
codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 11.9μs -> 9.34μs (27.6% faster)
# Count occurrences of describe and test in import
lines = result.split("\n")
import_line = next(line for line in lines if "import {" in line and "vitest" in line)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.To test or edit this optimization locally git merge codeflash/optimize-pr1266-2026-02-03T06.54.35
Click to see suggested changes
| needs_import = any(f"{global_name}(" in code or f"{global_name} (" in code for global_name in test_globals) | |
| if not needs_import: | |
| return code | |
| # Determine which globals are actually used in the code | |
| used_globals = [g for g in test_globals if f"{g}(" in code or f"{g} (" in code] | |
| if not used_globals: | |
| return code | |
| # Build the import statement | |
| import_statement = f"import {{ {', '.join(used_globals)} }} from 'vitest';\n" | |
| # Find the first line that isn't a comment or empty | |
| lines = code.split("\n") | |
| insert_index = 0 | |
| for i, line in enumerate(lines): | |
| stripped = line.strip() | |
| if stripped and not stripped.startswith("//") and not stripped.startswith("/*") and not stripped.startswith("*"): | |
| # Check if this line is an import/require - insert after imports | |
| if stripped.startswith("import ") or stripped.startswith("const ") or stripped.startswith("let "): | |
| continue | |
| insert_index = i | |
| break | |
| insert_index = i + 1 | |
| # Find the last import line to insert after it | |
| last_import_index = -1 | |
| for i, line in enumerate(lines): | |
| stripped = line.strip() | |
| if stripped.startswith("import ") and "from " in stripped: | |
| last_import_index = i | |
| # Compute used globals in a single pass to avoid scanning code twice | |
| used_globals = [g for g in test_globals if f"{g}(" in code or f"{g} (" in code] | |
| needs_import = bool(used_globals) | |
| if not needs_import: | |
| return code | |
| # Build the import statement | |
| import_statement = f"import {{ {', '.join(used_globals)} }} from 'vitest';\n" | |
| # Find the first line that isn't a comment or empty | |
| lines = code.split("\n") | |
| insert_index = 0 | |
| last_import_index = -1 | |
| for i, line in enumerate(lines): | |
| stripped = line.strip() | |
| if stripped.startswith("import ") and "from " in stripped: | |
| last_import_index = i | |
| if stripped and not stripped.startswith("//") and not stripped.startswith("/*") and not stripped.startswith("*"): | |
| # Check if this line is an import/require - insert after imports | |
| if stripped.startswith("import ") or stripped.startswith("const ") or stripped.startswith("let "): | |
| continue | |
| insert_index = i | |
| break | |
| insert_index = i + 1 | |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Too large diff
No description provided.