add async support in some AST visitors

[KRRT7]

PR Type

Enhancement, Tests

---

Description

• Add async support in AST analyses

• Extend unused helper detection to async

• Enhance global assignment collection logic

• Add comprehensive async-focused tests

---

Diagram Walkthrough

flowchart LR
  A["coverage_utils.extract_dependent_function"] -- "include AsyncFunctionDef" --> B["handles sync + async"]
  C["unused_definition_remover.detect_unused_helper_functions"] -- "locate entrypoint" --> D["supports async defs"]
  E["GlobalAssignmentCollector"] -- "tests ensure async-safety" --> F["collect only module-level"]
  G["OptimFunctionCollector"] -- "tests ensure async support" --> H["collects async/sync funcs"]

Loading

File Walkthrough

	Relevant files
Enhancement	coverage_utils.py Consider async defs in dependent function extraction          codeflash/code_utils/coverage_utils.py Include ast.AsyncFunctionDef in dependent lookup Ensure main function exclusion still applies +3/-1      unused_definition_remover.py Async-aware entrypoint detection in remover                            codeflash/context/unused_definition_remover.py Treat async and sync functions equally in entrypoint search Use isinstance with (FunctionDef, AsyncFunctionDef) +4/-1
Tests	test_code_context_extractor.py Tests for global assignment collection with async                tests/test_code_context_extractor.py Import GlobalAssignmentCollector for testing Add tests for async-aware global assignment collection Validate order and scoping with async/sync/classes +146/-1  test_code_replacement.py Tests for OptimFunctionCollector async coverage                    tests/test_code_replacement.py Import OptimFunctionCollector Add tests for collecting async/sync functions and methods Validate new functions and class methods detection +135/-0  test_code_utils.py Tests for async-aware dependent function extraction            tests/test_code_utils.py Import extract_dependent_function into tests Add tests for sync/async dependent extraction Cover edge and mixed async scenarios +81/-1    test_unused_helper_revert.py Async unused helper detection and revert tests                      tests/test_unused_helper_revert.py Import revert_unused_helper_functions Add async/sync mixed scenarios for unused helper detection Add revert tests for async helpers and aliasing Cover generators, coroutines, class async methods +614/-5

---

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit a6072b9)

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Possible Issue The set of dependent functions is built from only top-level defs in module body; nested functions or methods won't be considered. Confirm this is intended for dependency extraction when async support was added. dependent_functions = { node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) } Robustness When locating the entrypoint, both sync and async defs are considered, but selection relies on exact name match only. If multiple defs with same name exist across scopes (e.g., nested or class methods), this may pick the wrong node. Validate scope handling. if ( isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == function_to_optimize.function_name ): entrypoint_function_ast = node Test Stability Several tests write files and perform filesystem operations without using tmp_path fixtures consistently; ensure cleanup covers all cases and tests are isolated to avoid flakiness on parallel runs. 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 = """ 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 = """ 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 = """ 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 = """ 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_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 = """ 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) </details> </td></tr> </table>

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to a6072b9
Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Robustly collect module-level defs This only inspects top-level ast_tree.body and will miss functions nested inside classes or other scopes that may be relevant dependencies. If the intent is to capture any single dependent function at module scope, also ensure you filter only true module-level defs and avoid duplicates; otherwise, traverse with ast.walk to include async/sync defs consistently. codeflash/code_utils/coverage_utils.py [17-19] dependent_functions = { - node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + node.name + for node in ast.walk(ast_tree) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and isinstance(getattr(node, "parent", None), ast.Module) } Suggestion importance[1-10]: 4 __ Why: The existing code at lines 17-19 matches and was changed to include async defs, but it intentionally collects only top-level functions from ast_tree.body. The proposed change relies on .parent attributes that stdlib ast nodes don't have, making it incorrect as written.	Low
	Scope-limit entrypoint detection This will match any function with the same name anywhere, including class methods, leading to incorrect entrypoint identification for methods vs. module-level functions. Constrain the match to module-level nodes (or also match by class parent when optimizing a method). codeflash/context/unused_definition_remover.py [645-651] for node in ast.walk(optimized_ast): - if ( - isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and node.name == function_to_optimize.function_name - ): - entrypoint_function_ast = node - break + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == function_to_optimize.function_name: + # Ensure we're selecting the correct scope (module-level only) + if isinstance(getattr(node, "parent", None), ast.Module): + entrypoint_function_ast = node + break Suggestion importance[1-10]: 3 __ Why: The existing code at lines 645-651 matches; it currently finds any function by name. While scope-limiting can be helpful, the proposal again depends on .parent on ast nodes, which isn't available by default, so the suggestion is not directly applicable and risks being incorrect without a separate parent-annotation pass.	Low

---

Previous suggestions

Suggestions up to commit 8039ec4

Category	Suggestion	Impact
Possible issue	Avoid mismatching nested functions When matching entrypoints, include ast.AsyncFunctionDef as done, but also ensure parent context matches to avoid picking a same-named nested function. Verify the node is at the expected scope (e.g., class vs module) before assigning. codeflash/context/unused_definition_remover.py [645-649] -if ( - isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and node.name == function_to_optimize.function_name -): +if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == function_to_optimize.function_name: + # Ensure we only select functions at the desired scope (module or within specified parents) entrypoint_function_ast = node break Suggestion importance[1-10]: 5 __ Why: The suggestion correctly warns about ensuring scope when matching the entrypoint to avoid nested functions with the same name, which could improve correctness. However, the improved_code does not implement scope checking, so it's only a partial, moderate-impact suggestion.	Low
General	Handle async defs at module scope Ensure only top-level functions are considered; ast_tree.body can contain classes and other nodes but not nested defs. Also, guard against duplicate names by filtering only direct function definitions at module scope, which you already do; add a fallback to return False when no functions exist to avoid KeyError downstream. codeflash/code_utils/coverage_utils.py [17] -dependent_functions = {node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))} +dependent_functions = { + node.name + for node in ast_tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) +} Suggestion importance[1-10]: 2 __ Why: The existing code already filters top-level sync and async functions correctly using ast.FunctionDef and ast.AsyncFunctionDef. The improved_code is identical to existing_code and the content mentions a fallback not reflected in the code, so impact is minimal.	Low

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 14% (0.14x) speedup for CommentMapper.visit_FunctionDef in codeflash/code_utils/edit_generated_tests.py

⏱️ Runtime : 4.56 milliseconds → 4.01 milliseconds (best of 241 runs)

I created a new dependent PR with the suggested changes. Please review:

• ⚡️ Speed up method CommentMapper.visit_FunctionDef by 14% in PR #719 (async-support-for) #720

If you approve, it will be merged into this PR (branch async-support-for).

[github-actions]

Persistent review updated to latest commit a6072b9