Added a new module for function context and modified the main script to use it to get function dependencies.

codeflash/main.py

@@ -1,10 +1,11 @@
 from dataclasses import dataclass
 from typing import Optional
+import os
 
 from codeflash.discovery.functions_to_optimize import get_functions_to_optimize
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
 from codeflash.code_utils.code_extractor import get_code
-
+from codeflash.optimization.function_context import get_function_variables_definition
 
 @dataclass
 class args:
@@ -35,7 +36,8 @@ def main() -> None:
                     print("Could not find function to optimize")
                     continue
                 print(code_to_optimize)
-
+                function_dependencies = get_function_variables_definition(function_name, path, args.root)
+                print(function_dependencies)
 
     except Exception as e:
         print(e)

codeflash/optimization/function_context.py

@@ -0,0 +1,35 @@
+import os
+import jedi
+from codeflash.code_utils.code_extractor import get_code
+
+
+def belongs_to_function(name, function_name):
+    if name.full_name and name.full_name.startswith(name.module_name):
+        subname = name.full_name[len(name.module_name):]
+        if f'.{function_name}.' in subname:
+            return True
+    return False
+
+
+def get_function_variables_definition(function_name, file_path, project_root_path: str):
+    script = jedi.Script(path=file_path, project=jedi.Project(path=project_root_path))
+    sources = []
+    # The function name condition can be stricter so that it does not clash with other class names etc.
+    names = [ref for ref in script.get_names(all_scopes=True, definitions=False, references=True) if
+         ref.full_name and belongs_to_function(ref, function_name)]
+    for name in names:
+        definition = script.goto(line=name.line, column=name.column, follow_imports=True, follow_builtin_imports=False)
+        if definition:
+            definition_path = str(definition[0].module_path)
+            # The definition is part of this project and not defined within the original function
+            if definition_path.startswith(project_root_path + os.sep) and definition[0].full_name and not belongs_to_function(definition[0], function_name):
+                sources.append((name, definition[0], get_code(definition_path, definition[0].name)))
+    deduped_sources = []
+    existing_full_names = set()
+    for source in sources:
+        if source[0].full_name not in existing_full_names:
+            deduped_sources.append(source)
+            existing_full_names.add(source[0].full_name)
+    return deduped_sources
+
+

codeflash/verifier.py

@@ -1,6 +1,5 @@
 # Derived from https://github.com/openai/openai-cookbook/blob/main/examples/Unit_test_writing_using_a_multi-step_prompt.ipynb
 
-# imports needed to run the code in this notebook
 import ast  # used for detecting whether generated Python code is valid
 
 import openai  # used for calling the OpenAI API