implementation

Author: KRRT7

codeflash/code_utils/codeflash_wrap_decorator.py

@@ -5,14 +5,13 @@
 import os
 import pickle
 import time
-from collections.abc import Awaitable
 from functools import wraps
 from pathlib import Path
 from typing import Any, Callable, TypeVar
 
 from codeflash.code_utils.code_utils import get_run_tmp_file
 
-F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
+F = TypeVar("F", bound=Callable[..., Any])
 
 
 def extract_test_context_from_frame() -> tuple[str, str | None, str]:
@@ -33,9 +32,10 @@ def extract_test_context_from_frame() -> tuple[str, str | None, str]:
         del frame
 
 
-def codeflash_behavior_async(func: F) -> F:
+def codeflash_behavior_async(func):
+    """Async decorator for behavior analysis - collects timing data and function inputs/outputs."""
     function_name = func.__name__
-    line_id = f"{function_name}_{func.__code__.co_firstlineno}"
+    line_id = f"{func.__name__}_{func.__code__.co_firstlineno}"
     loop_index = int(os.environ.get("CODEFLASH_LOOP_INDEX", "1"))
 
     @wraps(func)
@@ -106,3 +106,63 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
         return return_value
 
     return async_wrapper
+
+
+def codeflash_performance_async(func):
+    """Async decorator for performance analysis - lightweight timing measurements only."""
+    function_name = func.__name__
+    line_id = f"{func.__name__}_{func.__code__.co_firstlineno}"
+    loop_index = int(os.environ.get("CODEFLASH_LOOP_INDEX", "1"))
+
+    @wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+        test_module_name, test_class_name, test_name = extract_test_context_from_frame()
+
+        test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{loop_index}"
+
+        if not hasattr(async_wrapper, "index"):
+            async_wrapper.index = {}
+        if test_id in async_wrapper.index:
+            async_wrapper.index[test_id] += 1
+        else:
+            async_wrapper.index[test_id] = 0
+
+        codeflash_test_index = async_wrapper.index[test_id]
+        invocation_id = f"{line_id}_{codeflash_test_index}"
+        test_stdout_tag = f"{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}"
+
+        print(f"!$######{test_stdout_tag}######$!")
+
+        exception = None
+        gc.disable()
+        try:
+            counter = time.perf_counter_ns()
+            ret = func(*args, **kwargs)
+
+            if inspect.isawaitable(ret):
+                counter = time.perf_counter_ns()
+                return_value = await ret
+            else:
+                return_value = ret
+
+            codeflash_duration = time.perf_counter_ns() - counter
+        except Exception as e:
+            codeflash_duration = time.perf_counter_ns() - counter
+            exception = e
+        finally:
+            gc.enable()
+
+        # For performance mode, include timing in the output tag like sync functions do
+        print(f"!######{test_stdout_tag}:{codeflash_duration}######!")
+
+        if exception:
+            raise exception
+        return return_value
+
+    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

@@ -322,13 +322,37 @@ 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,
+) -> 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:
+        return False, None
+
+
 def inject_profiling_into_existing_test(
     test_path: Path,
     call_positions: list[CodePosition],
     function_to_optimize: FunctionToOptimize,
     tests_project_root: Path,
     test_framework: str,
     mode: TestingMode = TestingMode.BEHAVIOR,
+    source_module_path: Path | None = None,
 ) -> tuple[bool, str | None]:
     with test_path.open(encoding="utf8") as f:
         test_code = f.read()
@@ -343,11 +367,10 @@ def inject_profiling_into_existing_test(
     import_visitor.visit(tree)
     func = import_visitor.imported_as
 
-    if func.is_async:
-        modified_code, decorator_added = add_async_decorator_to_function(test_code, func)
-        if decorator_added:
-            logger.debug(f"Applied @codeflash_trace_async decorator to async function {func.qualified_name}")
-            return True, modified_code
+    if func.is_async and source_module_path and source_module_path.exists():
+        source_success, instrumented_source = instrument_source_module_with_async_decorators(
+            source_module_path, func, mode
+        )
 
     tree = InjectPerfOnly(func, test_module_path, test_framework, call_positions, mode=mode).visit(tree)
     new_imports = [
@@ -739,21 +762,28 @@ def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR) -> ast.Fun
 
 
 class AsyncDecoratorAdder(cst.CSTTransformer):
-    """Transformer that adds @codeflash_trace_async decorator to async function definitions."""
+    """Transformer that adds async decorator to async function definitions."""
 
-    def __init__(self, function: FunctionToOptimize) -> None:
+    def __init__(self, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR) -> None:
         """Initialize the transformer.
 
         Args:
         ----
             function: The FunctionToOptimize object representing the target async function.
+            mode: The testing mode to determine which decorator to apply.
 
         """
         super().__init__()
         self.function = function
+        self.mode = mode
         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"
+        )
 
     def visit_ClassDef(self, node: cst.ClassDef) -> None:
         # Track when we enter a class
@@ -781,7 +811,7 @@ def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.Fu
 
             # Only add the decorator if it's not already there
             if not has_decorator:
-                new_decorator = cst.Decorator(decorator=cst.Name(value="codeflash_trace_async"))
+                new_decorator = cst.Decorator(decorator=cst.Name(value=self.decorator_name))
 
                 # Add our new decorator to the existing decorators
                 updated_decorators = [new_decorator, *list(updated_node.decorators)]
@@ -795,16 +825,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 == "codeflash_trace_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 == "codeflash_trace_async"
+            return decorator_node.func.value in {"codeflash_trace_async", "codeflash_behavior_async", "codeflash_performance_async"}
         return False
 
 
 class AsyncDecoratorImportAdder(cst.CSTTransformer):
-    """Transformer that adds the import for codeflash_trace_async."""
+    """Transformer that adds the import for async decorators."""
 
-    def __init__(self) -> None:
+    def __init__(self, mode: TestingMode = TestingMode.BEHAVIOR) -> None:
+        self.mode = mode
         self.has_import = False
 
     def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
@@ -819,48 +850,65 @@ def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
         ):
             # 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 == "codeflash_trace_async":
+                    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
         if self.has_import:
             return updated_node
 
+        # Choose import based on mode
+        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("from codeflash.code_utils.codeflash_wrap_decorator import codeflash_trace_async")
+        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) -> tuple[str, bool]:
-    """Add @codeflash_trace_async decorator to an async function definition.
+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:
     ----
         source_code: The source code to modify.
         function: The FunctionToOptimize object representing the target async function.
+        mode: The testing mode to determine which decorator to apply.
 
     Returns:
     -------
         Tuple of (modified_source_code, was_decorator_added).
 
     """
+    if not function.is_async:
+        return source_code, False
+        
     try:
         module = cst.parse_module(source_code)
 
         # Add the decorator to the function
-        decorator_transformer = AsyncDecoratorAdder(function)
+        decorator_transformer = AsyncDecoratorAdder(function, mode)
         module = module.visit(decorator_transformer)
 
         # Add the import if decorator was added
         if decorator_transformer.added_decorator:
-            import_transformer = AsyncDecoratorImportAdder()
+            import_transformer = AsyncDecoratorImportAdder(mode)
             module = module.visit(import_transformer)
 
         return isort.code(module.code, float_to_top=True), decorator_transformer.added_decorator
     except Exception as e:
         logger.exception(f"Error adding async decorator to function {function.qualified_name}: {e}")
         return source_code, False
+
+
+def create_instrumented_source_module_path(source_path: Path, temp_dir: Path) -> Path:
+    instrumented_filename = f"instrumented_{source_path.name}"
+    return temp_dir / instrumented_filename

codeflash/discovery/functions_to_optimize.py

@@ -159,7 +159,11 @@ def __str__(self) -> str:
 
     @property
     def qualified_name(self) -> str:
-        return self.function_name if self.parents == [] else f"{self.parents[0].name}.{self.function_name}"
+        if not self.parents:
+            return self.function_name
+        # Join all parent names with dots to handle nested classes properly
+        parent_path = ".".join(parent.name for parent in self.parents)
+        return f"{parent_path}.{self.function_name}"
 
     def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str:
         return f"{module_name_from_file_path(self.file_path, project_root_path)}.{self.qualified_name}"

pyproject.toml

@@ -43,6 +43,7 @@ dependencies = [
     "platformdirs>=4.3.7",
     "pygls>=1.3.1",
     "codeflash-benchmark",
+    "pytest-asyncio>=1.1.0",
 ]
 
 [project.urls]