linting / cleanup

Author: KRRT7

codeflash/code_utils/codeflash_wrap_decorator.py

@@ -32,7 +32,7 @@ def extract_test_context_from_frame() -> tuple[str, str | None, str]:
         del frame
 
 
-def codeflash_behavior_async(func):
+def codeflash_behavior_async(func: F) -> F:
     """Async decorator for behavior analysis - collects timing data and function inputs/outputs."""
     function_name = func.__name__
     line_id = f"{func.__name__}_{func.__code__.co_firstlineno}"
@@ -108,7 +108,7 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
     return async_wrapper
 
 
-def codeflash_performance_async(func):
+def codeflash_performance_async(func: F) -> F:
     """Async decorator for performance analysis - lightweight timing measurements only."""
     function_name = func.__name__
     line_id = f"{func.__name__}_{func.__code__.co_firstlineno}"
@@ -161,8 +161,3 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
 
     return async_wrapper
 
-
-# Convenience function to get the appropriate decorator based on mode
-def codeflash_trace_async(func):
-    """Default async decorator - uses behavior analysis mode."""
-    return codeflash_behavior_async(func)

codeflash/code_utils/instrument_existing_tests.py

@@ -323,25 +323,23 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
 
 
 def instrument_source_module_with_async_decorators(
-    source_path: Path,
-    function_to_optimize: FunctionToOptimize, 
-    mode: TestingMode = TestingMode.BEHAVIOR,
+    source_path: Path, function_to_optimize: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
 ) -> tuple[bool, str | None]:
     if not function_to_optimize.is_async:
         return False, None
-        
+
     try:
         with source_path.open(encoding="utf8") as f:
             source_code = f.read()
-            
+
         modified_code, decorator_added = add_async_decorator_to_function(source_code, function_to_optimize, mode)
-        
+
         if decorator_added:
             return True, modified_code
-        else:
-            return False, None
-            
-    except Exception as e:
+
+    except Exception:
+        return False, None
+    else:
         return False, None
 
 
@@ -361,7 +359,7 @@ def inject_profiling_into_existing_test(
     except SyntaxError:
         logger.exception(f"Syntax error in code in file - {test_path}")
         return False, None
-    
+
     test_module_path = module_name_from_file_path(test_path, tests_project_root)
     import_visitor = FunctionImportedAsVisitor(function_to_optimize)
     import_visitor.visit(tree)
@@ -779,7 +777,7 @@ def __init__(self, function: FunctionToOptimize, mode: TestingMode = TestingMode
         self.qualified_name_parts = function.qualified_name.split(".")
         self.context_stack = []
         self.added_decorator = False
-        
+
         # Choose decorator based on mode
         self.decorator_name = (
             "codeflash_behavior_async" if mode == TestingMode.BEHAVIOR else "codeflash_performance_async"
@@ -798,12 +796,9 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
         # Track when we enter a function
         self.context_stack.append(node.name.value)
 
-    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:  # noqa: ARG002
+    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
         # Check if this is an async function and matches our target
-        if (
-            original_node.asynchronous is not None
-            and self.context_stack == self.qualified_name_parts
-        ):
+        if original_node.asynchronous is not None and self.context_stack == self.qualified_name_parts:
             # Check if the decorator is already present
             has_decorator = any(
                 self._is_target_decorator(decorator.decorator) for decorator in original_node.decorators
@@ -825,9 +820,17 @@ def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.Fu
     def _is_target_decorator(self, decorator_node: cst.Name | cst.Attribute | cst.Call) -> bool:
         """Check if a decorator matches our target decorator name."""
         if isinstance(decorator_node, cst.Name):
-            return decorator_node.value in {"codeflash_trace_async", "codeflash_behavior_async", "codeflash_performance_async"}
+            return decorator_node.value in {
+                "codeflash_trace_async",
+                "codeflash_behavior_async",
+                "codeflash_performance_async",
+            }
         if isinstance(decorator_node, cst.Call) and isinstance(decorator_node.func, cst.Name):
-            return decorator_node.func.value in {"codeflash_trace_async", "codeflash_behavior_async", "codeflash_performance_async"}
+            return decorator_node.func.value in {
+                "codeflash_trace_async",
+                "codeflash_behavior_async",
+                "codeflash_performance_async",
+            }
         return False
 
 
@@ -847,15 +850,14 @@ def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
             and node.module.value.value.value == "codeflash"
             and node.module.value.attr.value == "code_utils"
             and node.module.attr.value == "codeflash_wrap_decorator"
+            and not isinstance(node.names, cst.ImportStar)
         ):
-            # Handle both ImportAlias and ImportStar
-            if not isinstance(node.names, cst.ImportStar):
-                decorator_name = (
-                    "codeflash_behavior_async" if self.mode == TestingMode.BEHAVIOR else "codeflash_performance_async"
-                )
-                for import_alias in node.names:
-                    if import_alias.name.value == decorator_name:
-                        self.has_import = True
+            decorator_name = (
+                "codeflash_behavior_async" if self.mode == TestingMode.BEHAVIOR else "codeflash_performance_async"
+            )
+            for import_alias in node.names:
+                if import_alias.name.value == decorator_name:
+                    self.has_import = True
 
     def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
         # If the import is already there, don't add it again
@@ -866,15 +868,17 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
         decorator_name = (
             "codeflash_behavior_async" if self.mode == TestingMode.BEHAVIOR else "codeflash_performance_async"
         )
-        
+
         # Parse the import statement into a CST node
         import_node = cst.parse_statement(f"from codeflash.code_utils.codeflash_wrap_decorator import {decorator_name}")
 
         # Add the import to the module's body
         return updated_node.with_changes(body=[import_node, *list(updated_node.body)])
 
 
-def add_async_decorator_to_function(source_code: str, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR) -> tuple[str, bool]:
+def add_async_decorator_to_function(
+    source_code: str, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
+) -> tuple[str, bool]:
     """Add async decorator to an async function definition.
 
     Args:
@@ -890,7 +894,7 @@ def add_async_decorator_to_function(source_code: str, function: FunctionToOptimi
     """
     if not function.is_async:
         return source_code, False
-        
+
     try:
         module = cst.parse_module(source_code)