Merge pull request #105 from codeflash-ai/helper_functions_fix_for_nested_stuff

prevent nested classes/ functions from showing up in code context

Author: alvin-r

codeflash/context/code_context_extractor.py

@@ -410,10 +410,12 @@ def get_function_sources_from_jedi(
                         and definition.full_name
                         and definition.type == "function"
                         and not belongs_to_function_qualified(definition, qualified_function_name)
+                        # Avoid nested functions or classes. Only class.function is allowed
+                        and len((qualified_name := get_qualified_name(definition.module_name, definition.full_name)).split(".")) <= 2
                     ):
                         function_source = FunctionSource(
                             file_path=definition_path,
-                            qualified_name=get_qualified_name(definition.module_name, definition.full_name),
+                            qualified_name=qualified_name,
                             fully_qualified_name=definition.full_name,
                             only_function_name=definition.name,
                             source_code=definition.get_line_code(),

tests/test_code_context_extractor.py

@@ -23,6 +23,12 @@ def innocent_bystander(self):
     def helper_method(self):
         return self.name
 
+    class NestedClass:
+        def __init__(self, name):
+            self.name = name
+
+        def nested_method(self):
+            return self.name
 
 def main_method():
     return "hello"
@@ -33,6 +39,7 @@ def __init__(self, name):
         self.name = name
 
     def main_method(self):
+        self.name = HelperClass.NestedClass("test").nested_method()
         return HelperClass(self.name).helper_method()
 
 
@@ -73,6 +80,8 @@ def test_code_replacement10() -> None:
     )
 
     code_ctx = get_code_optimization_context(function_to_optimize=func_top_optimize, project_root_path=file_path.parent)
+    qualified_names = {func.qualified_name for func in code_ctx.helper_functions}
+    assert qualified_names == {"HelperClass.helper_method"} # Nested method should not be in here
     read_write_context, read_only_context = code_ctx.read_writable_code, code_ctx.read_only_context_code
 
     expected_read_write_context = """
@@ -91,6 +100,7 @@ def __init__(self, name):
         self.name = name
 
     def main_method(self):
+        self.name = HelperClass.NestedClass("test").nested_method()
         return HelperClass(self.name).helper_method()
 """
     expected_read_only_context = """