mirror of
https://github.com/codeflash-ai/codeflash.git
synced 2026-05-04 18:25:17 +00:00
547c02e8bc
Move code_context_extractor.py and unused_definition_remover.py from codeflash/context/ to codeflash/languages/python/context/ and update all import sites.
2182 lines
74 KiB
Python
2182 lines
74 KiB
Python
"""Tests for unused helper function revert functionality."""
|
|
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
|
|
from codeflash.languages.python.context.unused_definition_remover import (
|
|
detect_unused_helper_functions,
|
|
revert_unused_helper_functions,
|
|
)
|
|
from codeflash.models.models import CodeStringsMarkdown
|
|
from codeflash.optimization.function_optimizer import FunctionOptimizer
|
|
from codeflash.verification.verification_utils import TestConfig
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_project():
|
|
"""Create a temporary project with test files."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
# Main file with function that calls helpers
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
def entrypoint_function(n):
|
|
\"\"\"Function that calls two helper functions.\"\"\"
|
|
result1 = helper_function_1(n)
|
|
result2 = helper_function_2(n)
|
|
return result1 + result2
|
|
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function.\"\"\"
|
|
return x * 3
|
|
""")
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
yield temp_dir, main_file, test_cfg
|
|
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_detect_unused_helper_functions(temp_project):
|
|
"""Test that unused helper functions are correctly detected."""
|
|
temp_dir, main_file, test_cfg = temp_project
|
|
|
|
# Optimized version that only calls one helper
|
|
optimized_code = """
|
|
```python:main.py
|
|
def entrypoint_function(n):
|
|
\"\"\"Optimized function that only calls one helper.\"\"\"
|
|
result1 = helper_function_1(n)
|
|
return result1 + n * 3 # Inlined helper_function_2
|
|
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function - MODIFIED VERSION should be reverted.\"\"\"
|
|
return x * 4 # This change should be reverted to original x * 3
|
|
```
|
|
"""
|
|
|
|
# Create FunctionToOptimize instance
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="entrypoint_function", qualified_name="entrypoint_function", parents=[]
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context to find helper functions
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect helper_function_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"helper_function_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# First modify the optimized code to include a MODIFIED unused helper
|
|
optimized_code_with_modified_helper = """
|
|
```python:main.py
|
|
def entrypoint_function(n):
|
|
\"\"\"Optimized function that only calls one helper.\"\"\"
|
|
result1 = helper_function_1(n)
|
|
return result1 + n * 3 # Inlined helper_function_2
|
|
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function - MODIFIED VERSION should be reverted.\"\"\"
|
|
return x * 7 # This should be reverted to x * 3
|
|
```
|
|
"""
|
|
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code_with_modified_helper), original_helper_code
|
|
)
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The entrypoint should be optimized
|
|
assert "result1 + n * 3" in final_content, "Entrypoint function should be optimized"
|
|
|
|
# helper_function_2 should be reverted to original (x * 3, NOT the modified x * 7)
|
|
assert "return x * 3" in final_content, "helper_function_2 should be reverted to original"
|
|
assert "return x * 7" not in final_content, "helper_function_2 should NOT contain the modified version"
|
|
|
|
# helper_function_1 should remain (it's still called)
|
|
assert "def helper_function_1(x):" in final_content, "helper_function_1 should still exist"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The entrypoint should be optimized
|
|
assert "result1 + n * 3" in final_content, "Entrypoint function should be optimized"
|
|
|
|
# helper_function_2 should be reverted to original (return x * 3, NOT the modified x * 4)
|
|
assert "return x * 3" in final_content, "helper_function_2 should be reverted to original"
|
|
assert "return x * 4" not in final_content, "helper_function_2 should NOT contain the modified version"
|
|
|
|
# helper_function_1 should remain as optimized (it's still called)
|
|
assert "def helper_function_1(x):" in final_content, "helper_function_1 should still exist"
|
|
|
|
|
|
def test_revert_unused_helper_functions(temp_project):
|
|
"""Test that unused helper functions are correctly reverted to original definitions."""
|
|
temp_dir, main_file, test_cfg = temp_project
|
|
|
|
# Optimized version that only calls one helper and modifies the unused one
|
|
optimized_code = """
|
|
```python:main.py
|
|
def entrypoint_function(n):
|
|
\"\"\"Optimized function that only calls one helper.\"\"\"
|
|
result1 = helper_function_1(n)
|
|
return result1 + n * 3 # Inlined helper_function_2
|
|
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Modified helper function - should be reverted.\"\"\"
|
|
return x * 4 # This change should be reverted
|
|
```
|
|
"""
|
|
|
|
# Create FunctionToOptimize instance
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="entrypoint_function", qualified_name="entrypoint_function", parents=[]
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Store original helper code
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
original_content = main_file.read_text()
|
|
|
|
# Test the new functionality - this should:
|
|
# 1. Apply the optimization
|
|
# 2. Detect unused helpers
|
|
# 3. Revert unused helpers to original definitions
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The entrypoint should be optimized (inline the helper_function_2 call)
|
|
assert "result1 + n * 3" in final_content, "Entrypoint function should be optimized"
|
|
|
|
# helper_function_2 should be reverted to original (return x * 3, not x * 4)
|
|
assert "return x * 3" in final_content, "helper_function_2 should be reverted to original"
|
|
assert "return x * 4" not in final_content, "helper_function_2 should not contain the optimized version"
|
|
|
|
# helper_function_1 should remain as optimized (it's still called)
|
|
assert "def helper_function_1(x):" in final_content, "helper_function_1 should still exist"
|
|
|
|
|
|
def test_no_unused_helpers_no_revert(temp_project):
|
|
"""Test that when all helpers are still used, nothing is reverted."""
|
|
temp_dir, main_file, test_cfg = temp_project
|
|
|
|
# Store original content to verify nothing changes
|
|
original_content = main_file.read_text()
|
|
|
|
revert_unused_helper_functions(temp_dir, [], {})
|
|
|
|
# Verify the file content remains unchanged
|
|
assert main_file.read_text() == original_content, "File should remain unchanged when no helpers to revert"
|
|
|
|
# Optimized version that still calls both helpers
|
|
optimized_code = """
|
|
```python:main.py
|
|
def entrypoint_function(n):
|
|
\"\"\"Optimized function that still calls both helpers.\"\"\"
|
|
result1 = helper_function_1(n)
|
|
result2 = helper_function_2(n)
|
|
return result1 + result2 # Still using both
|
|
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function - optimized.\"\"\"
|
|
return x << 1 # Optimized to use bit shift
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function - optimized.\"\"\"
|
|
return x * 3
|
|
```
|
|
"""
|
|
|
|
# Create FunctionToOptimize instance
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="entrypoint_function", qualified_name="entrypoint_function", parents=[]
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Store original helper code
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
# Test detection - should find no unused helpers
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
assert len(unused_helpers) == 0, "No helpers should be detected as unused"
|
|
|
|
# Apply optimization
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check final file content - should contain the optimized versions
|
|
final_content = main_file.read_text()
|
|
|
|
# Both helpers should be optimized
|
|
assert "x << 1" in final_content, "helper_function_1 should be optimized to use bit shift"
|
|
assert "result1 + result2" in final_content, "Entrypoint should still call both helpers"
|
|
|
|
|
|
def test_detect_unused_in_multi_file_project():
|
|
"""Test detection of unused helpers across multiple files."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
from helpers import helper_function_1, helper_function_2
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Function that calls helpers from another file.\"\"\"
|
|
result1 = helper_function_1(n)
|
|
result2 = helper_function_2(n)
|
|
return result1 + result2
|
|
""")
|
|
|
|
# Helper file
|
|
helper_file = temp_dir / "helpers.py"
|
|
helper_file.write_text("""
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function.\"\"\"
|
|
return x * 3
|
|
|
|
def helper_function_1(y): # Duplicate name to test line 575
|
|
\"\"\"Overloaded helper function.\"\"\"
|
|
return y + 10
|
|
""")
|
|
|
|
# Optimized version that only calls one helper with aliased import
|
|
optimized_code = """
|
|
```python:main.py
|
|
from helpers import helper_function_1 as h1
|
|
import helpers as h_module
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Optimized function that only calls one helper with aliasing.\"\"\"
|
|
result1 = h1(n) # Using aliased import
|
|
# Inlined helper_function_2 functionality: n * 3
|
|
return result1 + n * 3 # Fully inlined helper_function_2
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="entrypoint_function", qualified_name="entrypoint_function", parents=[]
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect helper_function_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"helper_function_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# First, simulate modified helper in the helper file
|
|
helper_file.write_text("""
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function - MODIFIED VERSION.\"\"\"
|
|
return x * 9 # This should be reverted to x * 3
|
|
""")
|
|
|
|
# Store original helper code (before modification)
|
|
original_helper_code = {
|
|
main_file: """
|
|
from helpers import helper_function_1, helper_function_2
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Function that calls helpers from another file.\"\"\"
|
|
result1 = helper_function_1(n)
|
|
result2 = helper_function_2(n)
|
|
return result1 + result2
|
|
""",
|
|
helper_file: """
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function.\"\"\"
|
|
return x * 3
|
|
""",
|
|
}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
# Check main file content
|
|
main_content = main_file.read_text()
|
|
assert "result1 + n * 3" in main_content, "Entrypoint function should be optimized"
|
|
assert "from helpers import helper_function_1" in main_content, "Import should be updated"
|
|
|
|
# Check helper file content - helper_function_2 should be reverted to original
|
|
helper_content = helper_file.read_text()
|
|
assert "def helper_function_1(x):" in helper_content, "helper_function_1 should still exist"
|
|
assert "def helper_function_2(x):" in helper_content, "helper_function_2 should exist"
|
|
assert "return x * 3" in helper_content, "helper_function_2 should be reverted to original"
|
|
assert "return x * 9" not in helper_content, "helper_function_2 should NOT contain the modified version"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# First, simulate modified helper in the helper file
|
|
helper_file.write_text("""
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function - MODIFIED VERSION.\"\"\"
|
|
return x * 5 # This should be reverted to x * 3
|
|
""")
|
|
|
|
# Store original helper code (before modification)
|
|
original_helper_code = {
|
|
main_file: """
|
|
from helpers import helper_function_1, helper_function_2
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Function that calls helpers from another file.\"\"\"
|
|
result1 = helper_function_1(n)
|
|
result2 = helper_function_2(n)
|
|
return result1 + result2
|
|
""",
|
|
helper_file: """
|
|
def helper_function_1(x):
|
|
\"\"\"First helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_function_2(x):
|
|
\"\"\"Second helper function.\"\"\"
|
|
return x * 3
|
|
""",
|
|
}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check main file content
|
|
main_content = main_file.read_text()
|
|
assert "result1 + n * 3" in main_content, "Entrypoint function should be optimized"
|
|
assert "from helpers import helper_function_1" in main_content, "Import should be updated"
|
|
|
|
# Check helper file content - helper_function_2 should be reverted to original
|
|
helper_content = helper_file.read_text()
|
|
assert "def helper_function_1(x):" in helper_content, "helper_function_1 should still exist"
|
|
assert "def helper_function_2(x):" in helper_content, "helper_function_2 should exist"
|
|
assert "return x * 3" in helper_content, "helper_function_2 should be reverted to original"
|
|
assert "return x * 5" not in helper_content, "helper_function_2 should NOT contain the modified version"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_class_method_entrypoint_with_helper_methods():
|
|
"""Test unused helper detection when entrypoint is a class method that calls other methods."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with class containing methods
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
class Calculator:
|
|
def entrypoint_method(self, n):
|
|
\"\"\"Main method that calls helper methods.\"\"\"
|
|
result1 = self.helper_method_1(n)
|
|
result2 = self.helper_method_2(n)
|
|
return result1 + result2
|
|
|
|
def helper_method_1(self, x):
|
|
\"\"\"First helper method.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_method_2(self, x):
|
|
\"\"\"Second helper method.\"\"\"
|
|
return x * 3
|
|
""")
|
|
|
|
# Optimized version that only calls one helper method
|
|
optimized_code = """
|
|
```python:main.py
|
|
class Calculator:
|
|
def entrypoint_method(self, n):
|
|
\"\"\"Optimized method that only calls one helper.\"\"\"
|
|
result1 = self.helper_method_1(n)
|
|
return result1 + n * 3 # Inlined helper_method_2
|
|
|
|
def helper_method_1(self, x):
|
|
\"\"\"First helper method.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_method_2(self, x):
|
|
\"\"\"Second helper method - should be reverted.\"\"\"
|
|
return x * 4
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance for class method
|
|
from codeflash.models.models import FunctionParent
|
|
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file,
|
|
function_name="entrypoint_method",
|
|
qualified_name="Calculator.entrypoint_method",
|
|
parents=[FunctionParent(name="Calculator", type="ClassDef")],
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect Calculator.helper_method_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"Calculator.helper_method_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# Update optimized code to include a MODIFIED unused helper
|
|
optimized_code_with_modified_helper = """
|
|
```python:main.py
|
|
class Calculator:
|
|
def entrypoint_method(self, n):
|
|
\"\"\"Optimized method that only calls one helper.\"\"\"
|
|
result1 = self.helper_method_1(n)
|
|
return result1 + n * 3 # Inlined helper_method_2
|
|
|
|
def helper_method_1(self, x):
|
|
\"\"\"First helper method.\"\"\"
|
|
return x * 2
|
|
|
|
def helper_method_2(self, x):
|
|
\"\"\"Second helper method - MODIFIED VERSION should be reverted.\"\"\"
|
|
return x * 8 # This should be reverted to x * 3
|
|
```
|
|
"""
|
|
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context,
|
|
CodeStringsMarkdown.parse_markdown_code(optimized_code_with_modified_helper),
|
|
original_helper_code,
|
|
)
|
|
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The entrypoint method should be optimized
|
|
assert "result1 + n * 3" in final_content, "Entrypoint method should be optimized"
|
|
|
|
# helper_method_2 should be reverted to original (x * 3, NOT the modified x * 8)
|
|
assert "return x * 3" in final_content, "helper_method_2 should be reverted to original"
|
|
assert "return x * 8" not in final_content, "helper_method_2 should NOT contain the modified version"
|
|
|
|
# helper_method_1 should remain (it's still called)
|
|
assert "def helper_method_1(self, x):" in final_content, "helper_method_1 should still exist"
|
|
|
|
# Test reversion
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The entrypoint method should be optimized
|
|
assert "result1 + n * 3" in final_content, "Entrypoint method should be optimized"
|
|
|
|
# helper_method_2 should be reverted to original
|
|
assert "x * 3" in final_content, "helper_method_2 should still exist"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_class_method_calls_external_helper_functions():
|
|
"""Test when class method calls external helper functions."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with class method that calls external helpers
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
def external_helper_1(x):
|
|
\"\"\"External helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def external_helper_2(x):
|
|
\"\"\"External helper function.\"\"\"
|
|
return x * 3
|
|
|
|
class Processor:
|
|
def process_data(self, n):
|
|
\"\"\"Method that calls external helper functions.\"\"\"
|
|
result1 = external_helper_1(n)
|
|
result2 = external_helper_2(n)
|
|
return result1 + result2
|
|
""")
|
|
|
|
# Optimized version that only calls one external helper
|
|
optimized_code = """
|
|
```python:main.py
|
|
def external_helper_1(x):
|
|
\"\"\"External helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def external_helper_2(x):
|
|
\"\"\"External helper function - should be reverted.\"\"\"
|
|
return x * 3
|
|
|
|
class Processor:
|
|
def process_data(self, n):
|
|
\"\"\"Optimized method that only calls one helper.\"\"\"
|
|
result1 = external_helper_1(n)
|
|
return result1 + n * 3 # Inlined external_helper_2
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance for class method
|
|
from codeflash.models.models import FunctionParent
|
|
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file,
|
|
function_name="process_data",
|
|
qualified_name="Processor.process_data",
|
|
parents=[FunctionParent(name="Processor", type="ClassDef")],
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect external_helper_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"external_helper_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# Update optimized code to include a MODIFIED unused helper
|
|
optimized_code_with_modified_helper = """
|
|
```python:main.py
|
|
def external_helper_1(x):
|
|
\"\"\"External helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def external_helper_2(x):
|
|
\"\"\"External helper function - MODIFIED VERSION should be reverted.\"\"\"
|
|
return x * 11 # This should be reverted to x * 3
|
|
|
|
class Processor:
|
|
def process_data(self, n):
|
|
\"\"\"Optimized method that only calls one helper.\"\"\"
|
|
result1 = external_helper_1(n)
|
|
return result1 + n * 3 # Inlined external_helper_2
|
|
```
|
|
"""
|
|
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context,
|
|
CodeStringsMarkdown.parse_markdown_code(optimized_code_with_modified_helper),
|
|
original_helper_code,
|
|
)
|
|
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The class method should be optimized
|
|
assert "result1 + n * 3" in final_content, "Process method should be optimized"
|
|
|
|
# external_helper_2 should be reverted to original (x * 3, NOT the modified x * 11)
|
|
assert "return x * 3" in final_content, "external_helper_2 should be reverted to original"
|
|
assert "return x * 11" not in final_content, "external_helper_2 should NOT contain the modified version"
|
|
|
|
# external_helper_1 should remain (it's still called)
|
|
assert "def external_helper_1(x):" in final_content, "external_helper_1 should still exist"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# Update optimized code to include a MODIFIED unused helper
|
|
optimized_code_with_modified_helper = """
|
|
```python:main.py
|
|
def external_helper_1(x):
|
|
\"\"\"External helper function.\"\"\"
|
|
return x * 2
|
|
|
|
def external_helper_2(x):
|
|
\"\"\"External helper function - MODIFIED VERSION should be reverted.\"\"\"
|
|
return x * 7 # This should be reverted to x * 3
|
|
|
|
class Processor:
|
|
def process_data(self, n):
|
|
\"\"\"Optimized method that only calls one helper.\"\"\"
|
|
result1 = external_helper_1(n)
|
|
return result1 + n * 3 # Inlined external_helper_2
|
|
```
|
|
"""
|
|
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context,
|
|
CodeStringsMarkdown.parse_markdown_code(optimized_code_with_modified_helper),
|
|
original_helper_code,
|
|
)
|
|
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The class method should be optimized
|
|
assert "result1 + n * 3" in final_content, "Process method should be optimized"
|
|
|
|
# external_helper_2 should be reverted to original (x * 3, NOT the modified x * 7)
|
|
assert "return x * 3" in final_content, "external_helper_2 should be reverted to original"
|
|
assert "return x * 7" not in final_content, "external_helper_2 should NOT contain the modified version"
|
|
|
|
# external_helper_1 should remain (it's still called)
|
|
assert "def external_helper_1(x):" in final_content, "external_helper_1 should still exist"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_nested_class_method_optimization():
|
|
"""Test optimization of methods in nested classes."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with nested class
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
def global_helper_1(x):
|
|
return x * 2
|
|
|
|
def global_helper_2(x):
|
|
return x * 3
|
|
|
|
class OuterClass:
|
|
class InnerProcessor:
|
|
def compute(self, n):
|
|
\"\"\"Method that calls global helper functions.\"\"\"
|
|
result1 = global_helper_1(n)
|
|
result2 = global_helper_2(n)
|
|
return result1 + result2
|
|
|
|
def local_helper(self, x):
|
|
return x + 1
|
|
""")
|
|
|
|
# Optimized version that inlines one helper
|
|
optimized_code = """
|
|
```python:main.py
|
|
def global_helper_1(x):
|
|
return x * 2
|
|
|
|
def global_helper_2(x):
|
|
return x * 3
|
|
|
|
class OuterClass:
|
|
class InnerProcessor:
|
|
def compute(self, n):
|
|
\"\"\"Optimized method.\"\"\"
|
|
result1 = global_helper_1(n)
|
|
return result1 + n * 3 # Inlined global_helper_2
|
|
|
|
def local_helper(self, x):
|
|
return x + 1
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Note: In practice, codeflash might not handle deeply nested classes,
|
|
# but we test the detection logic anyway
|
|
from codeflash.models.models import FunctionParent
|
|
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file,
|
|
function_name="compute",
|
|
qualified_name="OuterClass.InnerProcessor.compute",
|
|
parents=[
|
|
FunctionParent(name="OuterClass", type="ClassDef"),
|
|
FunctionParent(name="InnerProcessor", type="ClassDef"),
|
|
],
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Test detection directly (context extraction might not work for nested classes)
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize,
|
|
# Create a minimal context for testing
|
|
type(
|
|
"MockContext",
|
|
(),
|
|
{
|
|
"helper_functions": [
|
|
type(
|
|
"MockHelper",
|
|
(),
|
|
{
|
|
"qualified_name": "global_helper_1",
|
|
"only_function_name": "global_helper_1",
|
|
"fully_qualified_name": "main.global_helper_1",
|
|
"file_path": main_file,
|
|
"jedi_definition": type("MockJedi", (), {"type": "function"})(),
|
|
},
|
|
)(),
|
|
type(
|
|
"MockHelper",
|
|
(),
|
|
{
|
|
"qualified_name": "global_helper_2",
|
|
"only_function_name": "global_helper_2",
|
|
"fully_qualified_name": "main.global_helper_2",
|
|
"file_path": main_file,
|
|
"jedi_definition": type("MockJedi", (), {"type": "function"})(),
|
|
},
|
|
)(),
|
|
]
|
|
},
|
|
)(),
|
|
CodeStringsMarkdown.parse_markdown_code(optimized_code),
|
|
)
|
|
|
|
# Should detect global_helper_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"global_helper_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
# For nested class tests, we'll skip the complete workflow test since nested classes
|
|
# may not be fully supported by the optimizer, but we've verified detection works
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# Since this test uses nested classes which might not be fully supported,
|
|
# we'll only test with the mock context for detection but skip the full workflow test
|
|
# The other tests cover the complete workflow comprehensively
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_multi_file_import_styles():
|
|
"""Test detection with different import styles in multi-file projects."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
import utils
|
|
from math_helpers import add, multiply
|
|
from processors import process_data as pd
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Function using different import styles.\"\"\"
|
|
result1 = utils.compute(n) # Module.function style
|
|
result2 = add(n, 5) # Direct import style
|
|
result3 = multiply(n, 2) # Direct import style
|
|
result4 = pd(n) # Aliased import style
|
|
return result1 + result2 + result3 + result4
|
|
""")
|
|
|
|
# Utils file
|
|
utils_file = temp_dir / "utils.py"
|
|
utils_file.write_text("""
|
|
def compute(x):
|
|
\"\"\"Utility compute function.\"\"\"
|
|
return x * 10
|
|
|
|
def unused_util(x):
|
|
\"\"\"This utility function should be unused.\"\"\"
|
|
return x + 100
|
|
""")
|
|
|
|
# Math helpers file
|
|
math_file = temp_dir / "math_helpers.py"
|
|
math_file.write_text("""
|
|
def add(x, y):
|
|
\"\"\"Add two numbers.\"\"\"
|
|
return x + y
|
|
|
|
def multiply(x, y):
|
|
\"\"\"Multiply two numbers.\"\"\"
|
|
return x * y
|
|
|
|
def subtract(x, y):
|
|
\"\"\"Subtract function - should be unused.\"\"\"
|
|
return x - y
|
|
""")
|
|
|
|
# Processors file
|
|
processors_file = temp_dir / "processors.py"
|
|
processors_file.write_text("""
|
|
def process_data(x):
|
|
\"\"\"Process data.\"\"\"
|
|
return x ** 2
|
|
|
|
def clean_data(x):
|
|
\"\"\"Clean data - should be unused.\"\"\"
|
|
return x
|
|
""")
|
|
|
|
# Optimized version that only uses some functions
|
|
optimized_code = """
|
|
```python:main.py
|
|
import utils
|
|
from math_helpers import add
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Optimized function using fewer helpers.\"\"\"
|
|
result1 = utils.compute(n) # Still using utils.compute
|
|
result2 = add(n, 5) # Still using add
|
|
# Inlined multiply: result3 = n * 2
|
|
# Inlined process_data: result4 = n ** 2
|
|
return result1 + result2 + (n * 2) + (n ** 2)
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="entrypoint_function", qualified_name="entrypoint_function", parents=[]
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect multiply, process_data as unused (at minimum)
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
|
|
# The exact unused functions may vary based on what helpers are discovered by Jedi
|
|
# At minimum, we expect multiply to be detected as unused since it's not imported
|
|
assert "multiply" in unused_names, "Expected multiply to be detected as unused"
|
|
assert "process_data" in unused_names, "Expected process_data to be detected as unused"
|
|
assert "subtract" not in unused_names, "Expected subtract not to be detected as unused"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# First modify some helper files to simulate optimization changes
|
|
math_file.write_text("""
|
|
def add(x, y):
|
|
\"\"\"Add two numbers.\"\"\"
|
|
return x + y
|
|
|
|
def multiply(x, y):
|
|
\"\"\"Multiply two numbers - MODIFIED VERSION.\"\"\"
|
|
return x * y * 2 # This should be reverted to x * y
|
|
|
|
def subtract(x, y):
|
|
\"\"\"Subtract function - should be unused.\"\"\"
|
|
return x - y
|
|
""")
|
|
|
|
# Store original helper code
|
|
original_helper_code = {
|
|
main_file: """
|
|
import utils
|
|
from math_helpers import add, multiply
|
|
from processors import process_data as pd
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Function using different import styles.\"\"\"
|
|
result1 = utils.compute(n) # Module.function style
|
|
result2 = add(n, 5) # Direct import style
|
|
result3 = multiply(n, 2) # Direct import style
|
|
result4 = pd(n) # Aliased import style
|
|
return result1 + result2 + result3 + result4
|
|
""",
|
|
utils_file: utils_file.read_text(),
|
|
math_file: """
|
|
def add(x, y):
|
|
\"\"\"Add two numbers.\"\"\"
|
|
return x + y
|
|
|
|
def multiply(x, y):
|
|
\"\"\"Multiply two numbers.\"\"\"
|
|
return x * y
|
|
|
|
def subtract(x, y):
|
|
\"\"\"Subtract function - should be unused.\"\"\"
|
|
return x - y
|
|
""",
|
|
processors_file: processors_file.read_text(),
|
|
}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check main file content
|
|
main_content = main_file.read_text()
|
|
assert "(n * 2) + (n ** 2)" in main_content, "Entrypoint function should be optimized with inlined calculations"
|
|
assert "from math_helpers import add" in main_content, (
|
|
"Imports should be updated to only include used functions"
|
|
)
|
|
|
|
# Verify that unused helper files are reverted if they contained unused functions that were modified
|
|
math_content = math_file.read_text()
|
|
assert "def add(x, y):" in math_content, "add function should still exist"
|
|
# If multiply was unused and modified, it should be reverted
|
|
if "multiply" in unused_names:
|
|
assert "return x * y" in math_content, "multiply should be reverted to original if it was unused"
|
|
assert "return x * y * 2" not in math_content, (
|
|
"multiply should NOT contain the modified version if it was unused"
|
|
)
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_module_dot_function_import_style():
|
|
"""Test detection when helpers are called via module.function style."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
import calculator
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Function using module.function import style.\"\"\"
|
|
result1 = calculator.add_numbers(n, 10)
|
|
result2 = calculator.multiply_numbers(n, 5)
|
|
return result1 + result2
|
|
""")
|
|
|
|
# Calculator file
|
|
calc_file = temp_dir / "calculator.py"
|
|
calc_file.write_text("""
|
|
def add_numbers(x, y):
|
|
\"\"\"Add two numbers.\"\"\"
|
|
return x + y
|
|
|
|
def multiply_numbers(x, y):
|
|
\"\"\"Multiply two numbers.\"\"\"
|
|
return x * y
|
|
|
|
def divide_numbers(x, y):
|
|
\"\"\"Divide function - should be unused.\"\"\"
|
|
return x / y
|
|
""")
|
|
|
|
# Optimized version that only uses add_numbers
|
|
optimized_code = """
|
|
```python:main.py
|
|
import calculator
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Optimized function that inlines multiply.\"\"\"
|
|
result1 = calculator.add_numbers(n, 10)
|
|
# Inlined: result2 = n * 5
|
|
return result1 + (n * 5)
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="entrypoint_function", qualified_name="entrypoint_function", parents=[]
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect multiply_numbers and divide_numbers as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
|
|
# Check that multiply_numbers is detected as unused
|
|
assert "multiply_numbers" in unused_names, f"Expected 'multiply_numbers' to be unused, got: {unused_names}"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# First modify the calculator file to simulate optimization changes
|
|
calc_file.write_text("""
|
|
def add_numbers(x, y):
|
|
\"\"\"Add two numbers.\"\"\"
|
|
return x + y
|
|
|
|
def multiply_numbers(x, y):
|
|
\"\"\"Multiply two numbers - MODIFIED VERSION.\"\"\"
|
|
return x * y * 5 # This should be reverted to x * y
|
|
|
|
def divide_numbers(x, y):
|
|
\"\"\"Divide function - should be unused.\"\"\"
|
|
return x / y
|
|
""")
|
|
|
|
# Store original helper code
|
|
original_helper_code = {
|
|
main_file: """
|
|
import calculator
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Function using module.function import style.\"\"\"
|
|
result1 = calculator.add_numbers(n, 10)
|
|
result2 = calculator.multiply_numbers(n, 5)
|
|
return result1 + result2
|
|
""",
|
|
calc_file: """
|
|
def add_numbers(x, y):
|
|
\"\"\"Add two numbers.\"\"\"
|
|
return x + y
|
|
|
|
def multiply_numbers(x, y):
|
|
\"\"\"Multiply two numbers.\"\"\"
|
|
return x * y
|
|
|
|
def divide_numbers(x, y):
|
|
\"\"\"Divide function - should be unused.\"\"\"
|
|
return x / y
|
|
""",
|
|
}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check main file content
|
|
main_content = main_file.read_text()
|
|
assert "+ (n * 5)" in main_content, "Entrypoint function should be optimized with inlined multiplication"
|
|
assert "import calculator" in main_content, "Calculator import should remain"
|
|
|
|
# Check calculator file content - unused functions should be reverted if modified
|
|
calc_content = calc_file.read_text()
|
|
assert "def add_numbers(x, y):" in calc_content, "add_numbers should still exist"
|
|
assert "def multiply_numbers(x, y):" in calc_content, "multiply_numbers should exist"
|
|
assert "def divide_numbers(x, y):" in calc_content, "divide_numbers should remain as original"
|
|
# multiply_numbers should be reverted to original since it's unused
|
|
assert "return x * y" in calc_content, "multiply_numbers should be reverted to original"
|
|
assert "return x * y * 5" not in calc_content, "multiply_numbers should NOT contain the modified version"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# First modify the calculator file to simulate optimization changes
|
|
calc_file.write_text("""
|
|
def add_numbers(x, y):
|
|
\"\"\"Add two numbers.\"\"\"
|
|
return x + y
|
|
|
|
def multiply_numbers(x, y):
|
|
\"\"\"Multiply two numbers - MODIFIED VERSION.\"\"\"
|
|
return x * y * 3 # This should be reverted to x * y
|
|
|
|
def divide_numbers(x, y):
|
|
\"\"\"Divide function - should be unused.\"\"\"
|
|
return x / y
|
|
""")
|
|
|
|
# Store original helper code
|
|
original_helper_code = {
|
|
main_file: """
|
|
import calculator
|
|
|
|
def entrypoint_function(n):
|
|
\"\"\"Function using module.function import style.\"\"\"
|
|
result1 = calculator.add_numbers(n, 10)
|
|
result2 = calculator.multiply_numbers(n, 5)
|
|
return result1 + result2
|
|
""",
|
|
calc_file: """
|
|
def add_numbers(x, y):
|
|
\"\"\"Add two numbers.\"\"\"
|
|
return x + y
|
|
|
|
def multiply_numbers(x, y):
|
|
\"\"\"Multiply two numbers.\"\"\"
|
|
return x * y
|
|
|
|
def divide_numbers(x, y):
|
|
\"\"\"Divide function - should be unused.\"\"\"
|
|
return x / y
|
|
""",
|
|
}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check main file content
|
|
main_content = main_file.read_text()
|
|
assert "+ (n * 5)" in main_content, "Entrypoint function should be optimized with inlined multiplication"
|
|
assert "import calculator" in main_content, "Calculator import should remain"
|
|
|
|
# Check calculator file content - unused functions should be reverted if modified
|
|
calc_content = calc_file.read_text()
|
|
assert "def add_numbers(x, y):" in calc_content, "add_numbers should still exist"
|
|
assert "def multiply_numbers(x, y):" in calc_content, "multiply_numbers should exist"
|
|
assert "def divide_numbers(x, y):" in calc_content, "divide_numbers should remain as original"
|
|
# multiply_numbers should be reverted to original since it's unused
|
|
assert "return x * y" in calc_content, "multiply_numbers should be reverted to original"
|
|
assert "return x * y * 3" not in calc_content, "multiply_numbers should NOT contain the modified version"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_static_method_and_class_method():
|
|
"""Test optimization of static methods and class methods."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with static and class methods
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
def utility_function_1(x):
|
|
return x * 2
|
|
|
|
def utility_function_2(x):
|
|
return x * 3
|
|
|
|
class MathUtils:
|
|
@staticmethod
|
|
def calculate_static(n):
|
|
\"\"\"Static method that calls utility functions.\"\"\"
|
|
result1 = utility_function_1(n)
|
|
result2 = utility_function_2(n)
|
|
return result1 + result2
|
|
|
|
@classmethod
|
|
def calculate_class(cls, n):
|
|
\"\"\"Class method that calls utility functions.\"\"\"
|
|
result1 = utility_function_1(n)
|
|
result2 = utility_function_2(n)
|
|
return result1 - result2
|
|
""")
|
|
|
|
# Optimized static method that inlines one utility
|
|
optimized_static_code = """
|
|
```python:main.py
|
|
def utility_function_1(x):
|
|
return x * 2
|
|
|
|
def utility_function_2(x):
|
|
return x * 3
|
|
|
|
class MathUtils:
|
|
@staticmethod
|
|
def calculate_static(n):
|
|
\"\"\"Optimized static method.\"\"\"
|
|
result1 = utility_function_1(n)
|
|
return result1 + n * 3 # Inlined utility_function_2
|
|
|
|
@classmethod
|
|
def calculate_class(cls, n):
|
|
\"\"\"Class method that calls utility functions.\"\"\"
|
|
result1 = utility_function_1(n)
|
|
result2 = utility_function_2(n)
|
|
return result1 - result2
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Test static method optimization
|
|
from codeflash.models.models import FunctionParent
|
|
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file,
|
|
function_name="calculate_static",
|
|
qualified_name="MathUtils.calculate_static",
|
|
parents=[FunctionParent(name="MathUtils", type="ClassDef")],
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection for static method
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_static_code)
|
|
)
|
|
|
|
# Should detect utility_function_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"utility_function_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
|
|
# Update optimized code to include a MODIFIED unused helper
|
|
optimized_static_code_with_modified_helper = """
|
|
```python:main.py
|
|
def utility_function_1(x):
|
|
return x * 2
|
|
|
|
def utility_function_2(x):
|
|
return x * 6 # MODIFIED VERSION - should be reverted to x * 3
|
|
|
|
class MathUtils:
|
|
@staticmethod
|
|
def calculate_static(n):
|
|
\"\"\"Optimized static method.\"\"\"
|
|
result1 = utility_function_1(n)
|
|
return result1 + n * 3 # Inlined utility_function_2
|
|
|
|
@classmethod
|
|
def calculate_class(cls, n):
|
|
\"\"\"Class method that calls utility functions.\"\"\"
|
|
result1 = utility_function_1(n)
|
|
result2 = utility_function_2(n)
|
|
return result1 - result2
|
|
```
|
|
"""
|
|
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context,
|
|
CodeStringsMarkdown.parse_markdown_code(optimized_static_code_with_modified_helper),
|
|
original_helper_code,
|
|
)
|
|
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The static method should be optimized
|
|
assert "result1 + n * 3" in final_content, "Static method should be optimized"
|
|
|
|
# utility_function_2 should be reverted to original (x * 3, NOT the modified x * 6)
|
|
assert "return x * 3" in final_content, "utility_function_2 should be reverted to original"
|
|
assert "return x * 6" not in final_content, "utility_function_2 should NOT contain the modified version"
|
|
|
|
# utility_function_1 should remain (it's still called)
|
|
assert "def utility_function_1(x):" in final_content, "utility_function_1 should still exist"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_async_entrypoint_with_async_helpers():
|
|
"""Test that unused async helper functions are correctly detected when entrypoint is async."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with async entrypoint and async helpers
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
async def async_helper_1(x):
|
|
\"\"\"First async helper function.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_2(x):
|
|
\"\"\"Second async helper function.\"\"\"
|
|
return x * 3
|
|
|
|
async def async_entrypoint(n):
|
|
\"\"\"Async entrypoint function that calls async helpers.\"\"\"
|
|
result1 = await async_helper_1(n)
|
|
result2 = await async_helper_2(n)
|
|
return result1 + result2
|
|
""")
|
|
|
|
# Optimized version that only calls one async helper
|
|
optimized_code = """
|
|
```python:main.py
|
|
async def async_helper_1(x):
|
|
\"\"\"First async helper function.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_2(x):
|
|
\"\"\"Second async helper function - should be unused.\"\"\"
|
|
return x * 3
|
|
|
|
async def async_entrypoint(n):
|
|
\"\"\"Optimized async entrypoint that only calls one helper.\"\"\"
|
|
result1 = await async_helper_1(n)
|
|
return result1 + n * 3 # Inlined async_helper_2
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance for async function
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="async_entrypoint", parents=[], is_async=True
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect async_helper_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"async_helper_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_sync_entrypoint_with_async_helpers():
|
|
"""Test that unused async helper functions are detected when entrypoint is sync."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with sync entrypoint and async helpers
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
import asyncio
|
|
|
|
async def async_helper_1(x):
|
|
\"\"\"First async helper function.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_2(x):
|
|
\"\"\"Second async helper function.\"\"\"
|
|
return x * 3
|
|
|
|
def sync_entrypoint(n):
|
|
\"\"\"Sync entrypoint function that calls async helpers.\"\"\"
|
|
result1 = asyncio.run(async_helper_1(n))
|
|
result2 = asyncio.run(async_helper_2(n))
|
|
return result1 + result2
|
|
""")
|
|
|
|
# Optimized version that only calls one async helper
|
|
optimized_code = """
|
|
```python:main.py
|
|
import asyncio
|
|
|
|
async def async_helper_1(x):
|
|
\"\"\"First async helper function.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_2(x):
|
|
\"\"\"Second async helper function - should be unused.\"\"\"
|
|
return x * 3
|
|
|
|
def sync_entrypoint(n):
|
|
\"\"\"Optimized sync entrypoint that only calls one async helper.\"\"\"
|
|
result1 = asyncio.run(async_helper_1(n))
|
|
return result1 + n * 3 # Inlined async_helper_2
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance for sync function
|
|
function_to_optimize = FunctionToOptimize(file_path=main_file, function_name="sync_entrypoint", parents=[])
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect async_helper_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"async_helper_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_mixed_sync_and_async_helpers():
|
|
"""Test detection when both sync and async helpers are mixed."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with mixed sync and async helpers
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
import asyncio
|
|
|
|
def sync_helper_1(x):
|
|
\"\"\"Sync helper function.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_1(x):
|
|
\"\"\"Async helper function.\"\"\"
|
|
return x * 3
|
|
|
|
def sync_helper_2(x):
|
|
\"\"\"Another sync helper function.\"\"\"
|
|
return x * 4
|
|
|
|
async def async_helper_2(x):
|
|
\"\"\"Another async helper function.\"\"\"
|
|
return x * 5
|
|
|
|
async def mixed_entrypoint(n):
|
|
\"\"\"Async entrypoint function that calls both sync and async helpers.\"\"\"
|
|
sync_result = sync_helper_1(n)
|
|
async_result = await async_helper_1(n)
|
|
sync_result2 = sync_helper_2(n)
|
|
async_result2 = await async_helper_2(n)
|
|
return sync_result + async_result + sync_result2 + async_result2
|
|
""")
|
|
|
|
# Optimized version that only calls some helpers
|
|
optimized_code = """
|
|
```python:main.py
|
|
import asyncio
|
|
|
|
def sync_helper_1(x):
|
|
\"\"\"Sync helper function.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_1(x):
|
|
\"\"\"Async helper function.\"\"\"
|
|
return x * 3
|
|
|
|
def sync_helper_2(x):
|
|
\"\"\"Another sync helper function - should be unused.\"\"\"
|
|
return x * 4
|
|
|
|
async def async_helper_2(x):
|
|
\"\"\"Another async helper function - should be unused.\"\"\"
|
|
return x * 5
|
|
|
|
async def mixed_entrypoint(n):
|
|
\"\"\"Optimized async entrypoint that only calls some helpers.\"\"\"
|
|
sync_result = sync_helper_1(n)
|
|
async_result = await async_helper_1(n)
|
|
return sync_result + async_result + n * 4 + n * 5 # Inlined both helper_2 functions
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance for async function
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="mixed_entrypoint", parents=[], is_async=True
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect both sync_helper_2 and async_helper_2 as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"sync_helper_2", "async_helper_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_async_class_methods():
|
|
"""Test unused async method detection in classes."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with class containing async methods
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
class AsyncProcessor:
|
|
async def entrypoint_method(self, n):
|
|
\"\"\"Async main method that calls async helper methods.\"\"\"
|
|
result1 = await self.async_helper_method_1(n)
|
|
result2 = await self.async_helper_method_2(n)
|
|
return result1 + result2
|
|
|
|
async def async_helper_method_1(self, x):
|
|
\"\"\"First async helper method.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_method_2(self, x):
|
|
\"\"\"Second async helper method.\"\"\"
|
|
return x * 3
|
|
|
|
def sync_helper_method(self, x):
|
|
\"\"\"Sync helper method.\"\"\"
|
|
return x * 4
|
|
""")
|
|
|
|
# Optimized version that only calls one async helper
|
|
optimized_code = """
|
|
```python:main.py
|
|
class AsyncProcessor:
|
|
async def entrypoint_method(self, n):
|
|
\"\"\"Optimized async method that only calls one helper.\"\"\"
|
|
result1 = await self.async_helper_method_1(n)
|
|
return result1 + n * 3 # Inlined async_helper_method_2
|
|
|
|
async def async_helper_method_1(self, x):
|
|
\"\"\"First async helper method.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_method_2(self, x):
|
|
\"\"\"Second async helper method - should be unused.\"\"\"
|
|
return x * 3
|
|
|
|
def sync_helper_method(self, x):
|
|
\"\"\"Sync helper method - should be unused.\"\"\"
|
|
return x * 4
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance for async class method
|
|
from codeflash.models.models import FunctionParent
|
|
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file,
|
|
function_name="entrypoint_method",
|
|
parents=[FunctionParent(name="AsyncProcessor", type="ClassDef")],
|
|
is_async=True,
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect async_helper_method_2 as unused (sync_helper_method may not be discovered as helper)
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"AsyncProcessor.async_helper_method_2"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_async_helper_revert_functionality():
|
|
"""Test that unused async helper functions are correctly reverted to original definitions."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with async functions
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
async def async_helper_1(x):
|
|
\"\"\"First async helper function.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_2(x):
|
|
\"\"\"Second async helper function.\"\"\"
|
|
return x * 3
|
|
|
|
async def async_entrypoint(n):
|
|
\"\"\"Async entrypoint function that calls async helpers.\"\"\"
|
|
result1 = await async_helper_1(n)
|
|
result2 = await async_helper_2(n)
|
|
return result1 + result2
|
|
""")
|
|
|
|
# Optimized version that only calls one helper and modifies the unused one
|
|
optimized_code = """
|
|
```python:main.py
|
|
async def async_helper_1(x):
|
|
\"\"\"First async helper function.\"\"\"
|
|
return x * 2
|
|
|
|
async def async_helper_2(x):
|
|
\"\"\"Modified async helper function - should be reverted.\"\"\"
|
|
return x * 10 # This change should be reverted
|
|
|
|
async def async_entrypoint(n):
|
|
\"\"\"Optimized async entrypoint that only calls one helper.\"\"\"
|
|
result1 = await async_helper_1(n)
|
|
return result1 + n * 3 # Inlined async_helper_2
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance for async function
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="async_entrypoint", parents=[], is_async=True
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Store original helper code
|
|
original_helper_code = {main_file: main_file.read_text()}
|
|
|
|
# Apply optimization and test reversion
|
|
optimizer.replace_function_and_helpers_with_optimized_code(
|
|
code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code), original_helper_code
|
|
)
|
|
|
|
# Check final file content
|
|
final_content = main_file.read_text()
|
|
|
|
# The entrypoint should be optimized
|
|
assert "result1 + n * 3" in final_content, "Async entrypoint function should be optimized"
|
|
|
|
# async_helper_2 should be reverted to original (return x * 3, not x * 10)
|
|
assert "return x * 3" in final_content, "async_helper_2 should be reverted to original"
|
|
assert "return x * 10" not in final_content, "async_helper_2 should not contain the modified version"
|
|
|
|
# async_helper_1 should remain (it's still called)
|
|
assert "async def async_helper_1(x):" in final_content, "async_helper_1 should still exist"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_recursive_helper_function_not_detected_as_unused():
|
|
"""Test that recursive helper functions are NOT incorrectly detected as unused."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with recursive helper function
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
def gcd_recursive(a: int, b: int) -> int:
|
|
\"\"\"Calculate greatest common divisor using Euclidean algorithm with recursion.\"\"\"
|
|
if b == 0:
|
|
return a
|
|
return gcd_recursive(b, a % b)
|
|
""")
|
|
|
|
# Optimized version that still uses the recursive helper
|
|
optimized_code = """
|
|
```python:main.py
|
|
def gcd_recursive(a: int, b: int) -> int:
|
|
\"\"\"Calculate greatest common divisor using Euclidean algorithm with recursion.\"\"\"
|
|
if b == 0:
|
|
return a
|
|
return gcd_recursive(b, a % b)
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance
|
|
function_to_optimize = FunctionToOptimize(file_path=main_file, function_name="gcd_recursive", parents=[])
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should NOT detect gcd_recursive as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
|
|
assert "gcd_recursive" not in unused_names, (
|
|
f"Recursive function gcd_recursive should NOT be detected as unused, but got unused: {unused_names}"
|
|
)
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
|
|
def test_async_generators_and_coroutines():
|
|
"""Test detection with async generators and coroutines."""
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
|
|
try:
|
|
# Main file with async generators and coroutines
|
|
main_file = temp_dir / "main.py"
|
|
main_file.write_text("""
|
|
import asyncio
|
|
|
|
async def async_generator_helper(n):
|
|
\"\"\"Async generator helper.\"\"\"
|
|
for i in range(n):
|
|
yield i * 2
|
|
|
|
async def coroutine_helper(x):
|
|
\"\"\"Coroutine helper.\"\"\"
|
|
await asyncio.sleep(0.1)
|
|
return x * 3
|
|
|
|
async def another_coroutine_helper(x):
|
|
\"\"\"Another coroutine helper.\"\"\"
|
|
await asyncio.sleep(0.1)
|
|
return x * 4
|
|
|
|
async def async_entrypoint_with_generators(n):
|
|
\"\"\"Async entrypoint function that uses generators and coroutines.\"\"\"
|
|
results = []
|
|
async for value in async_generator_helper(n):
|
|
results.append(value)
|
|
|
|
final_result = await coroutine_helper(sum(results))
|
|
another_result = await another_coroutine_helper(n)
|
|
return final_result + another_result
|
|
""")
|
|
|
|
# Optimized version that doesn't use one of the coroutines
|
|
optimized_code = """
|
|
```python:main.py
|
|
import asyncio
|
|
|
|
async def async_generator_helper(n):
|
|
\"\"\"Async generator helper.\"\"\"
|
|
for i in range(n):
|
|
yield i * 2
|
|
|
|
async def coroutine_helper(x):
|
|
\"\"\"Coroutine helper.\"\"\"
|
|
await asyncio.sleep(0.1)
|
|
return x * 3
|
|
|
|
async def another_coroutine_helper(x):
|
|
\"\"\"Another coroutine helper - should be unused.\"\"\"
|
|
await asyncio.sleep(0.1)
|
|
return x * 4
|
|
|
|
async def async_entrypoint_with_generators(n):
|
|
\"\"\"Optimized async entrypoint that inlines one coroutine.\"\"\"
|
|
results = []
|
|
async for value in async_generator_helper(n):
|
|
results.append(value)
|
|
|
|
final_result = await coroutine_helper(sum(results))
|
|
return final_result + n * 4 # Inlined another_coroutine_helper
|
|
```
|
|
"""
|
|
|
|
# Create test config
|
|
test_cfg = TestConfig(
|
|
tests_root=temp_dir / "tests",
|
|
tests_project_rootdir=temp_dir,
|
|
project_root_path=temp_dir,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
)
|
|
|
|
# Create FunctionToOptimize instance for async function
|
|
function_to_optimize = FunctionToOptimize(
|
|
file_path=main_file, function_name="async_entrypoint_with_generators", parents=[], is_async=True
|
|
)
|
|
|
|
# Create function optimizer
|
|
optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=test_cfg,
|
|
function_to_optimize_source_code=main_file.read_text(),
|
|
)
|
|
|
|
# Get original code context
|
|
ctx_result = optimizer.get_code_optimization_context()
|
|
assert ctx_result.is_successful(), f"Failed to get context: {ctx_result.failure()}"
|
|
|
|
code_context = ctx_result.unwrap()
|
|
|
|
# Test unused helper detection
|
|
unused_helpers = detect_unused_helper_functions(
|
|
optimizer.function_to_optimize, code_context, CodeStringsMarkdown.parse_markdown_code(optimized_code)
|
|
)
|
|
|
|
# Should detect another_coroutine_helper as unused
|
|
unused_names = {uh.qualified_name for uh in unused_helpers}
|
|
expected_unused = {"another_coroutine_helper"}
|
|
|
|
assert unused_names == expected_unused, f"Expected unused: {expected_unused}, got: {unused_names}"
|
|
|
|
finally:
|
|
# Cleanup
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|