handled edge case for instrumenting codeflash trace

Author: alvin-r

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -12,29 +12,38 @@ def __init__(self, target_functions: set[tuple[str, str]]) -> None:
         self.target_functions = target_functions
         self.added_codeflash_trace = False
         self.class_name = ""
+        self.function_name = ""
         self.decorator = cst.Decorator(
             decorator=cst.Name(value="codeflash_trace")
         )
 
     def leave_ClassDef(self, original_node, updated_node):
-        self.class_name = ""
+        if self.class_name == original_node.name.value:
+            self.class_name = "" # Even if nested classes are not visited, this function is still called on them
         return updated_node
 
     def visit_ClassDef(self, node):
         if self.class_name: # Don't go into nested class
             return False
         self.class_name = node.name.value
 
+    def visit_FunctionDef(self, node):
+        if self.function_name: # Don't go into nested function
+            return False
+        self.function_name = node.name.value
+
     def leave_FunctionDef(self, original_node, updated_node):
+        if self.function_name == original_node.name.value:
+            self.function_name = ""
         if (self.class_name, original_node.name.value) in self.target_functions:
             # Add the new decorator after any existing decorators, so it gets executed first
             updated_decorators = list(updated_node.decorators) + [self.decorator]
             self.added_codeflash_trace = True
             return updated_node.with_changes(
                 decorators=updated_decorators
             )
-        else:
-            return updated_node
+
+        return updated_node
 
     def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
         # Create import statement for codeflash_trace

tests/test_instrument_codeflash_trace.py

@@ -466,4 +466,82 @@ def static_method_b():
 
         # Compare the modified content with expected content
         assert modified_content_1.strip() == expected_content_1.strip()
-        assert modified_content_2.strip() == expected_content_2.strip()
+        assert modified_content_2.strip() == expected_content_2.strip()
+
+
+def test_add_decorator_to_method_after_nested_class() -> None:
+    """Test adding decorator to a method that appears after a nested class definition."""
+    code = """
+class OuterClass:
+    class NestedClass:
+        def nested_method(self):
+            return "Hello from nested class method"
+
+    def target_method(self):
+        return "Hello from target method after nested class"
+"""
+
+    fto = FunctionToOptimize(
+        function_name="target_method",
+        file_path=Path("dummy_path.py"),
+        parents=[FunctionParent(name="OuterClass", type="ClassDef")]
+    )
+
+    modified_code = add_codeflash_decorator_to_code(
+        code=code,
+        functions_to_optimize=[fto]
+    )
+
+    expected_code = """
+from codeflash.benchmarking.codeflash_trace import codeflash_trace
+class OuterClass:
+    class NestedClass:
+        def nested_method(self):
+            return "Hello from nested class method"
+
+    @codeflash_trace
+    def target_method(self):
+        return "Hello from target method after nested class"
+"""
+
+    assert modified_code.strip() == expected_code.strip()
+
+
+def test_add_decorator_to_function_after_nested_function() -> None:
+    """Test adding decorator to a function that appears after a function with a nested function."""
+    code = """
+def function_with_nested():
+    def inner_function():
+        return "Hello from inner function"
+
+    return inner_function()
+
+def target_function():
+    return "Hello from target function after nested function"
+"""
+
+    fto = FunctionToOptimize(
+        function_name="target_function",
+        file_path=Path("dummy_path.py"),
+        parents=[]
+    )
+
+    modified_code = add_codeflash_decorator_to_code(
+        code=code,
+        functions_to_optimize=[fto]
+    )
+
+    expected_code = """
+from codeflash.benchmarking.codeflash_trace import codeflash_trace
+def function_with_nested():
+    def inner_function():
+        return "Hello from inner function"
+
+    return inner_function()
+
+@codeflash_trace
+def target_function():
+    return "Hello from target function after nested function"
+"""
+
+    assert modified_code.strip() == expected_code.strip()