cleaned up comments and removed unused old file

alvin-r · alvin-r · commit e9d6bd12a87e · 2025-03-11T16:15:07.000-07:00
diff --git a/codeflash/optimization/function_context.py b/codeflash/optimization/function_context.py
@@ -1,28 +1,10 @@
 from __future__ import annotations
 
-import ast
-import os
-import re
-from collections import defaultdict
-from typing import TYPE_CHECKING
-
-import jedi
-import tiktoken
 from jedi.api.classes import Name
-
-from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.code_extractor import get_code
 from codeflash.code_utils.code_utils import (
     get_qualified_name,
-    module_name_from_file_path,
-    path_belongs_to_site_packages,
-)
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-from codeflash.models.models import FunctionParent, FunctionSource
-
-if TYPE_CHECKING:
-    from pathlib import Path
 
+)
 
 def belongs_to_method(name: Name, class_name: str, method_name: str) -> bool:
     """Check if the given name belongs to the specified method."""
@@ -58,242 +40,4 @@ def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> b
                 return get_qualified_name(name.module_name, name.full_name) == qualified_function_name
         return False
     except ValueError:
-        return False
-
-#
-# def get_type_annotation_context(
-#     function: FunctionToOptimize, jedi_script: jedi.Script, project_root_path: Path
-# ) -> tuple[list[FunctionSource], set[tuple[str, str]]]:
-#     function_name: str = function.function_name
-#     file_path: Path = function.file_path
-#     file_contents: str = file_path.read_text(encoding="utf8")
-#     try:
-#         module: ast.Module = ast.parse(file_contents)
-#     except SyntaxError as e:
-#         logger.exception(f"get_type_annotation_context - Syntax error in code: {e}")
-#         return [], set()
-#     sources: list[FunctionSource] = []
-#     ast_parents: list[FunctionParent] = []
-#     contextual_dunder_methods = set()
-#
-#     def get_annotation_source(
-#         j_script: jedi.Script, name: str, node_parents: list[FunctionParent], line_no: int, col_no: str
-#     ) -> None:
-#         try:
-#             definition: list[Name] = j_script.goto(
-#                 line=line_no, column=col_no, follow_imports=True, follow_builtin_imports=False
-#             )
-#         except Exception as ex:
-#             if hasattr(name, "full_name"):
-#                 logger.exception(f"Error while getting definition for {name.full_name}: {ex}")
-#             else:
-#                 logger.exception(f"Error while getting definition: {ex}")
-#             definition = []
-#         if definition:  # TODO can be multiple definitions
-#             definition_path = definition[0].module_path
-#
-#             # The definition is part of this project and not defined within the original function
-#             if (
-#                 str(definition_path).startswith(str(project_root_path) + os.sep)
-#                 and definition[0].full_name
-#                 and not path_belongs_to_site_packages(definition_path)
-#                 and not belongs_to_function(definition[0], function_name)
-#             ):
-#                 source_code = get_code([FunctionToOptimize(definition[0].name, definition_path, node_parents[:-1])])
-#                 if source_code[0]:
-#                     sources.append(
-#                         FunctionSource(
-#                             fully_qualified_name=definition[0].full_name,
-#                             jedi_definition=definition[0],
-#                             source_code=source_code[0],
-#                             file_path=definition_path,
-#                             qualified_name=definition[0].full_name.removeprefix(definition[0].module_name + "."),
-#                             only_function_name=definition[0].name,
-#                         )
-#                     )
-#                     contextual_dunder_methods.update(source_code[1])
-#
-#     def visit_children(
-#         node: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef | ast.Module, node_parents: list[FunctionParent]
-#     ) -> None:
-#         child: ast.AST | ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef | ast.Module
-#         for child in ast.iter_child_nodes(node):
-#             visit(child, node_parents)
-#
-#     def visit_all_annotation_children(
-#         node: ast.Subscript | ast.Name | ast.BinOp, node_parents: list[FunctionParent]
-#     ) -> None:
-#         if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
-#             visit_all_annotation_children(node.left, node_parents)
-#             visit_all_annotation_children(node.right, node_parents)
-#         if isinstance(node, ast.Name) and hasattr(node, "id"):
-#             name: str = node.id
-#             line_no: int = node.lineno
-#             col_no: int = node.col_offset
-#             get_annotation_source(jedi_script, name, node_parents, line_no, col_no)
-#         if isinstance(node, ast.Subscript):
-#             if hasattr(node, "slice"):
-#                 if isinstance(node.slice, ast.Subscript):
-#                     visit_all_annotation_children(node.slice, node_parents)
-#                 elif isinstance(node.slice, ast.Tuple):
-#                     for elt in node.slice.elts:
-#                         if isinstance(elt, (ast.Name, ast.Subscript)):
-#                             visit_all_annotation_children(elt, node_parents)
-#                 elif isinstance(node.slice, ast.Name):
-#                     visit_all_annotation_children(node.slice, node_parents)
-#             if hasattr(node, "value"):
-#                 visit_all_annotation_children(node.value, node_parents)
-#
-#     def visit(
-#         node: ast.AST | ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef | ast.Module,
-#         node_parents: list[FunctionParent],
-#     ) -> None:
-#         if isinstance(node, (ast.Module, ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
-#             if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
-#                 if node.name == function_name and node_parents == function.parents:
-#                     arg: ast.arg
-#                     for arg in node.args.args:
-#                         if arg.annotation:
-#                             visit_all_annotation_children(arg.annotation, node_parents)
-#                     if node.returns:
-#                         visit_all_annotation_children(node.returns, node_parents)
-#
-#             if not isinstance(node, ast.Module):
-#                 node_parents.append(FunctionParent(node.name, type(node).__name__))
-#             visit_children(node, node_parents)
-#             if not isinstance(node, ast.Module):
-#                 node_parents.pop()
-#
-#     visit(module, ast_parents)
-#
-#     return sources, contextual_dunder_methods
-
-
-# def get_function_variables_definitions(
-#     function_to_optimize: FunctionToOptimize, project_root_path: Path
-# ) -> tuple[list[FunctionSource], set[tuple[str, str]]]:
-#     function_name = function_to_optimize.function_name
-#     file_path = function_to_optimize.file_path
-#     script = jedi.Script(path=file_path, project=jedi.Project(path=project_root_path))
-#     sources: list[FunctionSource] = []
-#     contextual_dunder_methods = set()
-#     # TODO: The function name condition can be stricter so that it does not clash with other class names etc.
-#     # TODO: The function could have been imported as some other name,
-#     #  we should be checking for the translation as well. Also check for the original function name.
-#     names = []
-#     for ref in script.get_names(all_scopes=True, definitions=False, references=True):
-#         if ref.full_name:
-#             if function_to_optimize.parents:
-#                 # Check if the reference belongs to the specified class when FunctionParent is provided
-#                 if belongs_to_method(ref, function_to_optimize.parents[-1].name, function_name):
-#                     names.append(ref)
-#             elif belongs_to_function(ref, function_name):
-#                 names.append(ref)
-#
-#     for name in names:
-#         try:
-#             definitions: list[Name] = name.goto(follow_imports=True, follow_builtin_imports=False)
-#         except Exception as e:
-#             try:
-#                 logger.exception(f"Error while getting definition for {name.full_name}: {e}")
-#             except Exception as e:
-#                 # name.full_name can also throw exceptions sometimes
-#                 logger.exception(f"Error while getting definition: {e}")
-#             definitions = []
-#         if definitions:
-#             # TODO: there can be multiple definitions, see how to handle such cases
-#             definition = definitions[0]
-#             definition_path = definition.module_path
-#
-#             # The definition is part of this project and not defined within the original function
-#             if (
-#                 str(definition_path).startswith(str(project_root_path) + os.sep)
-#                 and not path_belongs_to_site_packages(definition_path)
-#                 and definition.full_name
-#                 and not belongs_to_function(definition, function_name)
-#             ):
-#                 module_name = module_name_from_file_path(definition_path, project_root_path)
-#                 m = re.match(rf"{module_name}\.(.*)\.{definitions[0].name}", definitions[0].full_name)
-#                 parents = []
-#                 if m:
-#                     parents = [FunctionParent(m.group(1), "ClassDef")]
-#
-#                 source_code = get_code(
-#                     [FunctionToOptimize(function_name=definitions[0].name, file_path=definition_path, parents=parents)]
-#                 )
-#                 if source_code[0]:
-#                     sources.append(
-#                         FunctionSource(
-#                             fully_qualified_name=definition.full_name,
-#                             jedi_definition=definition,
-#                             source_code=source_code[0],
-#                             file_path=definition_path,
-#                             qualified_name=definition.full_name.removeprefix(definition.module_name + "."),
-#                             only_function_name=definition.name,
-#                         )
-#                     )
-#                     contextual_dunder_methods.update(source_code[1])
-#     annotation_sources, annotation_dunder_methods = get_type_annotation_context(
-#         function_to_optimize, script, project_root_path
-#     )
-#     sources[:0] = annotation_sources  # prepend the annotation sources
-#     contextual_dunder_methods.update(annotation_dunder_methods)
-#     existing_fully_qualified_names = set()
-#     no_parent_sources: dict[Path, dict[str, set[FunctionSource]]] = defaultdict(lambda: defaultdict(set))
-#     parent_sources = set()
-#     for source in sources:
-#         if (fully_qualified_name := source.fully_qualified_name) not in existing_fully_qualified_names:
-#             if not source.qualified_name.count("."):
-#                 no_parent_sources[source.file_path][source.qualified_name].add(source)
-#             else:
-#                 parent_sources.add(source)
-#             existing_fully_qualified_names.add(fully_qualified_name)
-#     deduped_parent_sources = [
-#         source
-#         for source in parent_sources
-#         if source.file_path not in no_parent_sources
-#         or source.qualified_name.rpartition(".")[0] not in no_parent_sources[source.file_path]
-#     ]
-#     deduped_no_parent_sources = [
-#         source for k1 in no_parent_sources for k2 in no_parent_sources[k1] for source in no_parent_sources[k1][k2]
-#     ]
-#     return deduped_no_parent_sources + deduped_parent_sources, contextual_dunder_methods
-#
-#
-# MAX_PROMPT_TOKENS = 4096  # 128000  # gpt-4-128k
-#
-#
-# def get_constrained_function_context_and_helper_functions(
-#     function_to_optimize: FunctionToOptimize,
-#     project_root_path: Path,
-#     code_to_optimize: str,
-#     max_tokens: int = MAX_PROMPT_TOKENS,
-# ) -> tuple[str, list[FunctionSource], set[tuple[str, str]]]:
-#     helper_functions, dunder_methods = get_function_variables_definitions(function_to_optimize, project_root_path)
-#     tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
-#     code_to_optimize_tokens = tokenizer.encode(code_to_optimize)
-#
-#     if not function_to_optimize.parents:
-#         helper_functions_sources = [function.source_code for function in helper_functions]
-#     else:
-#         helper_functions_sources = [
-#             function.source_code
-#             for function in helper_functions
-#             if not function.qualified_name.count(".")
-#             or function.qualified_name.split(".")[0] != function_to_optimize.parents[0].name
-#         ]
-#     helper_functions_tokens = [len(tokenizer.encode(function)) for function in helper_functions_sources]
-#
-#     context_list = []
-#     context_len = len(code_to_optimize_tokens)
-#     logger.debug(f"ORIGINAL CODE TOKENS LENGTH: {context_len}")
-#     logger.debug(f"ALL DEPENDENCIES TOKENS LENGTH: {sum(helper_functions_tokens)}")
-#     for function_source, source_len in zip(helper_functions_sources, helper_functions_tokens):
-#         if context_len + source_len <= max_tokens:
-#             context_list.append(function_source)
-#             context_len += source_len
-#         else:
-#             break
-#     logger.debug(f"FINAL OPTIMIZATION CONTEXT TOKENS LENGTH: {context_len}")
-#     helper_code: str = "\n".join(context_list)
-#     return helper_code, helper_functions, dunder_methods
+        return False
diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
@@ -48,7 +48,6 @@
     BestOptimization,
     CodeOptimizationContext,
     FunctionCalledInTest,
-    FunctionParent,
     GeneratedTests,
     GeneratedTestsList,
     OptimizationSet,
@@ -58,7 +57,6 @@
     TestFiles,
     TestingMode,
 )
-# from codeflash.optimization.function_context import get_constrained_function_context_and_helper_functions
 from codeflash.result.create_pr import check_create_pr, existing_tests_source_for
 from codeflash.result.critic import coverage_critic, performance_gain, quantity_of_tests_critic, speedup_critic
 from codeflash.result.explanation import Explanation
diff --git a/tests/test_function_dependencies.py b/tests/test_function_dependencies.py
@@ -6,7 +6,6 @@
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.either import is_successful
 from codeflash.models.models import FunctionParent
-# from codeflash.optimization.function_context import get_function_variables_definitions
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.verification_utils import TestConfig
 
diff --git a/tests/test_type_annotation_context.py b/tests/test_type_annotation_context.py
