Remove test fns with the benchmark fixture while instrumenting existing unit test files

[aseembits93]

PR Type

Enhancement, Tests

---

Description

Introduce BenchmarkFunctionRemover AST transformer.
Implement removal function for benchmark tests.
Integrate removal into test instrumentation.
Add extensive removal logic tests.

---

Changes walkthrough 📝

	Relevant files
Enhancement	code_replacer.py Add BenchmarkFunctionRemover transformer                                  codeflash/code_utils/code_replacer.py Added Union to typing imports. Introduced BenchmarkFunctionRemover transformer. Implemented benchmark detection and removal. Added remove_benchmark_functions util. +111/-1  instrument_existing_tests.py Integrate benchmark function removal                                          codeflash/code_utils/instrument_existing_tests.py Imported remove_benchmark_functions. Applied remover in test injection pipeline. +3/-0
Tests	test_code_replacement.py Add tests for benchmark removal                                                    tests/test_code_replacement.py Imported new remover classes/functions. Added TestBenchmarkFunctionRemover suite. Added TestRemoveBenchmarkFunctions suite. Added TestBenchmarkDetectionMethods tests. +377/-1

---

Need help?

Type /help how to ... in the comments thread for any questions about PR-Agent usage.

Check out the documentation for more information.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 6c1c2a4)

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 Docstring Mismatch The remove_benchmark_functions docstring claims to return a tuple of (modified_source_code, set_of_removed_function_names), but the implementation only returns an AST. Update the docstring or adjust the return value accordingly. def remove_benchmark_functions(tree: AST) -> AST: """Remove benchmark functions from Python source code. Args: tree: Python source code as ast module Returns: Tuple of (modified_source_code, set_of_removed_function_names) """ try: # Create and apply the transformer remover = BenchmarkFunctionRemover() return remover.visit(tree) except Exception as e: print(f"Error processing code: {e}") return tree Error Handling remove_benchmark_functions catches all exceptions and uses print() to report errors. It should use the project logger (e.g., logger.error) to surface issues in CI and avoid silent failures. try: # Create and apply the transformer remover = BenchmarkFunctionRemover() return remover.visit(tree) except Exception as e: print(f"Error processing code: {e}") return tree Empty Class Body visit_ClassDef may produce classes with an empty body after removing benchmark methods. Unparsing an empty class body can lead to invalid syntax. Consider adding a pass statement when the body is empty. def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: """Visit class definitions and remove benchmark methods.""" original_body = node.body[:] new_body = [] for item in original_body: if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)): if not self._uses_benchmark_fixture(item): new_body.append(self.visit(item)) else: new_body.append(self.visit(item)) node.body = new_body return node

[aseembits93]

todo : improve tests

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 6c1c2a4
Explore these optional code suggestions:

Category	Suggestion	Impact
General	Use logger.error instead of print Use the existing logger to report errors instead of print, so that error messages integrate with the application's logging configuration. codeflash/code_utils/code_replacer.py [117-133] def remove_benchmark_functions(tree: AST) -> AST: ... try: - # Create and apply the transformer remover = BenchmarkFunctionRemover() return remover.visit(tree) except Exception as e: - print(f"Error processing code: {e}") + logger.error("Error removing benchmark functions: %s", e) return tree Suggestion importance[1-10]: 6 __ Why: Switching from print to logger.error aligns with the project’s logging strategy and improves error reporting without altering core logic.	Low
	Detect benchmark in all arg types Also check keyword-only arguments, varargs, and kwargs for "benchmark" to catch all possible fixture injections. codeflash/code_utils/code_replacer.py [31-36] def _uses_benchmark_fixture(self, node: Union[ast.FunctionDef, ast.AsyncFunctionDef]) -> bool: """Check if a function uses the benchmark fixture.""" - # Check function arguments for 'benchmark' parameter - for arg in node.args.args: + # Check positional and keyword-only arguments + for arg in list(node.args.args) + list(node.args.kwonlyargs): if arg.arg == "benchmark": return True + # Check varargs and kwargs + if node.args.vararg and node.args.vararg.arg == "benchmark": + return True + if node.args.kwarg and node.args.kwarg.arg == "benchmark": + return True Suggestion importance[1-10]: 5 __ Why: Extending _uses_benchmark_fixture to cover keyword-only, varargs, and kwargs increases fixture detection accuracy in edge cases.	Low

---

Previous suggestions

Suggestions up to commit e353f38

Category	Suggestion	Impact
General	Use logger.exception and fix docstring Update the docstring to accurately reflect that this function returns an AST, not a tuple, and replace the print call with logger.exception to ensure errors are properly logged with tracebacks. codeflash/code_utils/code_replacer.py [117-133] def remove_benchmark_functions(tree: AST) -> AST: - """Remove benchmark functions from Python source code. + """Remove benchmark functions from an AST. Args: - tree: Python source code as ast module + tree: Python source code AST Returns: - Tuple of (modified_source_code, set_of_removed_function_names) - + AST with benchmark functions removed. """ try: - # Create and apply the transformer remover = BenchmarkFunctionRemover() return remover.visit(tree) - except Exception as e: - print(f"Error processing code: {e}") + logger.exception(f"Error removing benchmark functions: {e}") return tree Suggestion importance[1-10]: 6 __ Why: The suggestion corrects the docstring to match the actual return type and replaces the bare print with logger.exception to ensure errors are logged with full tracebacks, improving consistency and maintainability.	Low

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 58% (0.58x) speedup for BenchmarkFunctionRemover.visit_AsyncFunctionDef in codeflash/code_utils/code_replacer.py

⏱️ Runtime : 23.8 microseconds → 15.1 microseconds (best of 65 runs)

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

• ⚡️ Speed up method BenchmarkFunctionRemover.visit_AsyncFunctionDef by 58% in PR #313 (skip-benchmark-instrumentation) #314

If you approve, it will be merged into this PR (branch skip-benchmark-instrumentation).

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 61% (0.61x) speedup for BenchmarkFunctionRemover.visit_ClassDef in codeflash/code_utils/code_replacer.py

⏱️ Runtime : 44.6 microseconds → 27.6 microseconds (best of 60 runs)

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

• ⚡️ Speed up method BenchmarkFunctionRemover.visit_ClassDef by 61% in PR #313 (skip-benchmark-instrumentation) #315

If you approve, it will be merged into this PR (branch skip-benchmark-instrumentation).

[github-actions]

Persistent review updated to latest commit 6c1c2a4

[aseembits93]

closing for now as there are other ways of achieving the same goal