codeflash-ai/codeflash
1
0
Fork 0
mirror of https://github.com/codeflash-ai/codeflash.git synced 2026-05-04 18:25:17 +00:00
Code Issues 18 Projects Releases Packages Wiki Activity Actions

refactor: replace jedi_definition with definition_type on FunctionSource

Browse source 
Store only the type string instead of the full Jedi Name object,
removing the need for arbitrary_types_allowed and the runtime
dependency on jedi in the model layer.
This commit is contained in:
Kevin Turcios 2026-02-10 04:57:10 -05:00
parent 6428a22eba
commit a96918766f
8 changed files with 26 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
codeflash/code_utils/code_replacer.py
View file

@ -871,8 +871,7 @@ def replace_optimized_code(
[
callee.qualified_name
for callee in code_context.helper_functions
if callee.file_path == module_path
and (callee.jedi_definition is None or callee.jedi_definition.type != "class")
if callee.file_path == module_path and callee.definition_type != "class"
]
),
candidate.source_code,

8
codeflash/context/code_context_extractor.py
View file

@ -265,7 +265,7 @@ def get_code_optimization_context_for_language(
fully_qualified_name=helper.qualified_name,
only_function_name=helper.name,
source_code=helper.source_code,
jedi_definition=None,
definition_type=None,
)
)
@ -488,7 +488,7 @@ def get_function_to_optimize_as_function_source(
fully_qualified_name=name.full_name,
only_function_name=name.name,
source_code=name.get_line_code(),
jedi_definition=name,
definition_type=name.type,
)
except Exception as e:
logger.exception(f"Error while getting function source: {e}")
@ -544,7 +544,7 @@ def get_function_sources_from_jedi(
fully_qualified_name=definition.full_name,
only_function_name=definition.name,
source_code=definition.get_line_code(),
jedi_definition=definition,
definition_type=definition.type,
)
file_path_to_function_source[definition_path].add(function_source)
function_source_list.append(function_source)
@ -562,7 +562,7 @@ def get_function_sources_from_jedi(
fully_qualified_name=f"{definition.full_name}.__init__",
only_function_name="__init__",
source_code=definition.get_line_code(),
jedi_definition=definition,
definition_type=definition.type,
)
file_path_to_function_source[definition_path].add(function_source)
function_source_list.append(function_source)

8
codeflash/context/unused_definition_remover.py
View file

@ -630,8 +630,8 @@ def _analyze_imports_in_optimized_code(
helpers_by_file_and_func = defaultdict(dict)
helpers_by_file = defaultdict(list) # preserved for "import module"
for helper in code_context.helper_functions:
jedi_type = helper.jedi_definition.type if helper.jedi_definition else None
if jedi_type != "class": # Include when jedi_definition is None (non-Python)
jedi_type = helper.definition_type
if jedi_type != "class": # Include when definition_type is None (non-Python)
func_name = helper.only_function_name
module_name = helper.file_path.stem
# Cache function lookup for this (module, func)
@ -789,8 +789,8 @@ def detect_unused_helper_functions(
unused_helpers = []
entrypoint_file_path = function_to_optimize.file_path
for helper_function in code_context.helper_functions:
jedi_type = helper_function.jedi_definition.type if helper_function.jedi_definition else None
if jedi_type != "class": # Include when jedi_definition is None (non-Python)
jedi_type = helper_function.definition_type
if jedi_type != "class": # Include when definition_type is None (non-Python)
# Check if the helper function is called using multiple name variants
helper_qualified_name = helper_function.qualified_name
helper_simple_name = helper_function.only_function_name

5
codeflash/models/models.py
View file

@ -25,7 +25,6 @@ from pathlib import Path
from re import Pattern
from typing import Any, NamedTuple, Optional, cast
from jedi.api.classes import Name
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, ValidationError, model_validator
from pydantic.dataclasses import dataclass
@ -136,14 +135,14 @@ class CoverReturnCode(IntEnum):
ERROR = 2
@dataclass(frozen=True, config={"arbitrary_types_allowed": True})
@dataclass(frozen=True)
class FunctionSource:
file_path: Path
qualified_name: str
fully_qualified_name: str
only_function_name: str
source_code: str
jedi_definition: Name | None = None # None for non-Python languages
definition_type: str | None = None # e.g. "function", "class"; None for non-Python languages
def __eq__(self, other: object) -> bool:
if not isinstance(other, FunctionSource):

4
codeflash/optimization/function_optimizer.py
View file

@ -1501,8 +1501,8 @@ class FunctionOptimizer:
self.function_to_optimize.qualified_name
)
for helper_function in code_context.helper_functions:
# Skip class definitions (jedi_definition may be None for non-Python languages)
if helper_function.jedi_definition is None or helper_function.jedi_definition.type != "class":
# Skip class definitions (definition_type may be None for non-Python languages)
if helper_function.definition_type != "class":
read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name)
for module_abspath, qualified_names in read_writable_functions_by_file_path.items():
did_update |= replace_function_definitions_in_module(

34
tests/test_code_replacement.py
View file

@ -1,6 +1,5 @@
from __future__ import annotations
import dataclasses
import os
import re
from collections import defaultdict
@ -19,28 +18,13 @@ from codeflash.code_utils.code_replacer import (
replace_functions_in_file,
)
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, FunctionParent
from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, FunctionParent, FunctionSource
from codeflash.optimization.function_optimizer import FunctionOptimizer
from codeflash.verification.verification_utils import TestConfig
os.environ["CODEFLASH_API_KEY"] = "cf-test-key"
@dataclasses.dataclass
class JediDefinition:
type: str
@dataclasses.dataclass
class FakeFunctionSource:
file_path: Path
qualified_name: str
fully_qualified_name: str
only_function_name: str
source_code: str
jedi_definition: JediDefinition
class Args:
disable_imports_sorting = True
formatter_cmds = ["disabled"]
@ -1137,7 +1121,7 @@ class TestResults(BaseModel):
preexisting_objects = find_preexisting_objects(original_code)
helper_functions = [
FakeFunctionSource(
FunctionSource(
file_path=Path(
"/Users/saurabh/Library/CloudStorage/Dropbox/codeflash/cli/codeflash/verification/test_results.py"
),
@ -1145,7 +1129,7 @@ class TestResults(BaseModel):
fully_qualified_name="codeflash.verification.test_results.TestType",
only_function_name="TestType",
source_code="",
jedi_definition=JediDefinition(type="class"),
definition_type="class",
)
]
@ -1160,7 +1144,7 @@ class TestResults(BaseModel):
helper_functions_by_module_abspath = defaultdict(set)
for helper_function in helper_functions:
if helper_function.jedi_definition.type != "class":
if helper_function.definition_type != "class":
helper_functions_by_module_abspath[helper_function.file_path].add(helper_function.qualified_name)
for module_abspath, qualified_names in helper_functions_by_module_abspath.items():
new_code: str = replace_functions_and_add_imports(
@ -1352,21 +1336,21 @@ def cosine_similarity_top_k(
preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] = find_preexisting_objects(original_code)
helper_functions = [
FakeFunctionSource(
FunctionSource(
file_path=(Path(__file__).parent / "code_to_optimize" / "math_utils.py").resolve(),
qualified_name="Matrix",
fully_qualified_name="code_to_optimize.math_utils.Matrix",
only_function_name="Matrix",
source_code="",
jedi_definition=JediDefinition(type="class"),
definition_type="class",
),
FakeFunctionSource(
FunctionSource(
file_path=(Path(__file__).parent / "code_to_optimize" / "math_utils.py").resolve(),
qualified_name="cosine_similarity",
fully_qualified_name="code_to_optimize.math_utils.cosine_similarity",
only_function_name="cosine_similarity",
source_code="",
jedi_definition=JediDefinition(type="function"),
definition_type="function",
),
]
@ -1425,7 +1409,7 @@ def cosine_similarity_top_k(
)
helper_functions_by_module_abspath = defaultdict(set)
for helper_function in helper_functions:
if helper_function.jedi_definition.type != "class":
if helper_function.definition_type != "class":
helper_functions_by_module_abspath[helper_function.file_path].add(helper_function.qualified_name)
for module_abspath, qualified_names in helper_functions_by_module_abspath.items():
new_helper_code: str = replace_functions_and_add_imports(

5
tests/test_function_dependencies.py
View file

@ -151,10 +151,9 @@ def test_class_method_dependencies() -> None:
# The code_context above should have the topologicalSortUtil function in it
assert len(code_context.helper_functions) == 1
assert (
code_context.helper_functions[0].jedi_definition.full_name
== "test_function_dependencies.Graph.topologicalSortUtil"
code_context.helper_functions[0].fully_qualified_name == "test_function_dependencies.Graph.topologicalSortUtil"
)
assert code_context.helper_functions[0].jedi_definition.name == "topologicalSortUtil"
assert code_context.helper_functions[0].only_function_name == "topologicalSortUtil"
assert (
code_context.helper_functions[0].fully_qualified_name == "test_function_dependencies.Graph.topologicalSortUtil"
)

4
tests/test_unused_helper_revert.py
View file

@ -915,7 +915,7 @@ class OuterClass:
"only_function_name": "global_helper_1",
"fully_qualified_name": "main.global_helper_1",
"file_path": main_file,
"jedi_definition": type("MockJedi", (), {"type": "function"})(),
"definition_type": "function",
},
)(),
type(
@ -926,7 +926,7 @@ class OuterClass:
"only_function_name": "global_helper_2",
"fully_qualified_name": "main.global_helper_2",
"file_path": main_file,
"jedi_definition": type("MockJedi", (), {"type": "function"})(),
"definition_type": "function",
},
)(),
]