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

added helpers of helpers into readonly context, and refactored code slightly

Browse source
This commit is contained in:
Alvin Ryanputra 2024-12-31 17:19:09 -08:00
parent 641088fad5
commit a10b399dbe
6 changed files with 779 additions and 227 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
code_to_optimize/code_directories/retriever/transform_utils.py Normal file
View file

@ -0,0 +1,27 @@
from code_to_optimize.code_directories.retriever.utils import DataProcessor
class DataTransformer:
def __init__(self):
self.data = None
def transform(self, data):
self.data = data
return self.data
def transform_using_own_method(self, data):
return self.transform(data)
def transform_using_same_file_function(self, data):
return update_data(data)
def transform_data_all_same_file(self, data):
new_data = update_data(data)
return self.transform_using_own_method(new_data)
def circular_dependency(self, data):
return DataProcessor().circular_dependency(data)
def update_data(data):
return data + " updated"

20
code_to_optimize/code_directories/retriever/utils.py
View file

@ -1,5 +1,9 @@
import math
from transform_utils import DataTransformer
GLOBAL_VAR = 10
class DataProcessor:
"""A class for processing data."""
@ -25,3 +29,19 @@ class DataProcessor:
def do_something(self):
print("something")
def transform_data(self, data: str) -> str:
"""Transform the processed data"""
return DataTransformer().transform(data)
def transform_data_own_method(self, data: str) -> str:
"""Transform the processed data using own method"""
return DataTransformer().transform_using_own_method(data)
def transform_data_same_file_function(self, data: str) -> str:
"""Transform the processed data using a function from the same file"""
return DataTransformer().transform_using_same_file_function(data)
def circular_dependency(self, data: str) -> str:
"""Test circular dependency"""
return DataTransformer().circular_dependency(data)

311
codeflash/context/code_context_extractor.py
View file

@ -15,99 +15,36 @@ from codeflash.code_utils.code_extractor import add_needed_imports_from_module
from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.models.models import CodeString, CodeStringsMarkdown
from codeflash.optimization.function_context import belongs_to_class, belongs_to_function
from codeflash.optimization.function_context import belongs_to_function_qualified
def get_code_optimization_context(
function_to_optimize: FunctionToOptimize, project_root_path: Path, token_limit: int = 8000
) -> 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))
file_path_to_qualified_function_names = defaultdict(set)
file_path_to_qualified_function_names[file_path].add(function_to_optimize.qualified_name)
read_only_code_markdown = CodeStringsMarkdown()
final_read_writable_code = ""
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_class(ref, function_to_optimize.parents[-1].name) and belongs_to_function(
ref, function_name
):
names.append(ref)
elif belongs_to_function(ref, function_name):
names.append(ref)
# Get qualified names and fully qualified names(fqn) of helpers
helpers_of_fto, helpers_of_fto_fqn = get_file_path_to_helper_functions_dict(
{function_to_optimize.file_path: {function_to_optimize.qualified_name}}, project_root_path
)
helpers_of_helpers, helpers_of_helpers_fqn = get_file_path_to_helper_functions_dict(
helpers_of_fto, project_root_path
)
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
# Add function to optimize
helpers_of_fto[function_to_optimize.file_path].add(function_to_optimize.qualified_name)
helpers_of_fto_fqn[function_to_optimize.file_path].add(
function_to_optimize.qualified_name_with_modules_from_root(project_root_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)
and definition.module_name != definition.full_name
):
file_path_to_qualified_function_names[definition_path].add(
get_qualified_name(definition.module_name, definition.full_name)
)
for file_path, qualified_function_names in file_path_to_qualified_function_names.items():
try:
og_code_containing_helpers = file_path.read_text("utf8")
except Exception as e:
logger.exception(f"Error while parsing {file_path}: {e}")
continue
try:
read_writable_code = get_read_writable_code(og_code_containing_helpers, qualified_function_names)
except ValueError as e:
logger.debug(f"Error while getting read-writable code: {e}")
continue
if read_writable_code:
final_read_writable_code += f"\n{read_writable_code}"
final_read_writable_code = add_needed_imports_from_module(
src_module_code=og_code_containing_helpers,
dst_module_code=final_read_writable_code,
src_path=file_path,
dst_path=file_path,
project_root=project_root_path,
helper_functions_fqn=qualified_function_names,
)
try:
read_only_code = get_read_only_code(og_code_containing_helpers, qualified_function_names)
except ValueError as e:
logger.debug(f"Error while getting read-only code: {e}")
continue
read_only_code_with_imports = CodeString(
code=add_needed_imports_from_module(
src_module_code=og_code_containing_helpers,
dst_module_code=read_only_code,
src_path=file_path,
dst_path=file_path,
project_root=project_root_path,
helper_functions_fqn=qualified_function_names,
),
file_path=Path(file_path),
)
if read_only_code_with_imports.code:
read_only_code_markdown.code_strings.append(read_only_code_with_imports)
# Extract code
final_read_writable_code = get_all_read_writable_code(helpers_of_fto, helpers_of_fto_fqn, project_root_path)
read_only_code_markdown = get_all_read_only_code_context(
helpers_of_fto,
helpers_of_fto_fqn,
helpers_of_helpers,
helpers_of_helpers_fqn,
project_root_path,
remove_docstrings=False,
)
# Handle token limits
tokenizer = tiktoken.encoding_for_model("gpt-4o")
@ -121,12 +58,85 @@ def get_code_optimization_context(
return CodeString(code=final_read_writable_code).code, read_only_code_markdown.markdown
logger.debug("Code context has exceeded token limit, removing docstrings from read-only code")
# Get read-only code context again, this time without docstrings
# Extract read only code without docstrings
read_only_code_no_docstring_markdown = get_all_read_only_code_context(
helpers_of_fto,
helpers_of_fto_fqn,
helpers_of_helpers,
helpers_of_helpers_fqn,
project_root_path,
remove_docstrings=True,
)
read_only_code_no_docstring_markdown_tokens = len(tokenizer.encode(read_only_code_no_docstring_markdown.markdown))
total_tokens = final_read_writable_tokens + read_only_code_no_docstring_markdown_tokens
if total_tokens <= token_limit:
return CodeString(code=final_read_writable_code).code, read_only_code_no_docstring_markdown.markdown
logger.debug("Code context has exceeded token limit, removing read-only code")
return CodeString(code=final_read_writable_code).code, ""
def get_all_read_writable_code(
helpers_of_fto: dict[Path, set[str]], helpers_of_fto_fqn: dict[Path, set[str]], project_root_path: Path
) -> str:
final_read_writable_code = ""
# Extract code from file paths that contain fto and first degree helpers
for file_path, qualified_function_names in helpers_of_fto.items():
try:
original_code = file_path.read_text("utf8")
except Exception as e:
logger.exception(f"Error while parsing {file_path}: {e}")
continue
try:
read_writable_code = get_read_writable_code(original_code, qualified_function_names)
except ValueError as e:
logger.debug(f"Error while getting read-writable code: {e}")
continue
if read_writable_code:
final_read_writable_code += f"\n{read_writable_code}"
final_read_writable_code = add_needed_imports_from_module(
src_module_code=original_code,
dst_module_code=final_read_writable_code,
src_path=file_path,
dst_path=file_path,
project_root=project_root_path,
helper_functions_fqn=helpers_of_fto_fqn[file_path],
)
return final_read_writable_code
def get_all_read_only_code_context(
helpers_of_fto: dict[Path, set[str]],
helpers_of_fto_fqn: dict[Path, set[str]],
helpers_of_helpers: dict[Path, set[str]],
helpers_of_helpers_fqn: dict[Path, set[str]],
project_root_path: Path,
remove_docstrings: bool = False,
) -> CodeStringsMarkdown:
# Rearrange to remove overlaps, so we only access each file path once
helpers_of_helpers_no_overlap = defaultdict(set)
helpers_of_helpers_no_overlap_fqn = defaultdict(set)
for file_path in helpers_of_helpers:
if file_path in helpers_of_fto:
# Remove duplicates, in case a helper of helper is also a helper of fto
helpers_of_helpers[file_path] -= helpers_of_fto[file_path]
helpers_of_helpers_fqn[file_path] -= helpers_of_fto_fqn[file_path]
else:
helpers_of_helpers_no_overlap[file_path] = helpers_of_helpers[file_path]
helpers_of_helpers_no_overlap_fqn[file_path] = helpers_of_helpers_fqn[file_path]
read_only_code_markdown = CodeStringsMarkdown()
for file_path, qualified_function_names in file_path_to_qualified_function_names.items():
# Extract code from file paths that contain fto and first degree helpers. helpers of helpers may also be included if they are in the same files
for file_path, qualified_function_names in helpers_of_fto.items():
try:
original_code = file_path.read_text("utf8")
except Exception as e:
logger.exception(f"Error while parsing {file_path}: {e}")
continue
try:
read_only_code = get_read_only_code(
og_code_containing_helpers, qualified_function_names, remove_docstrings=True
original_code, qualified_function_names, helpers_of_helpers.get(file_path, set()), remove_docstrings
)
except ValueError as e:
logger.debug(f"Error while getting read-only code: {e}")
@ -134,24 +144,93 @@ def get_code_optimization_context(
read_only_code_with_imports = CodeString(
code=add_needed_imports_from_module(
src_module_code=og_code_containing_helpers,
src_module_code=original_code,
dst_module_code=read_only_code,
src_path=file_path,
dst_path=file_path,
project_root=project_root_path,
helper_functions_fqn=qualified_function_names,
helper_functions_fqn=helpers_of_fto_fqn[file_path] | helpers_of_helpers_fqn[file_path],
),
file_path=Path(file_path),
)
if read_only_code_with_imports.code:
read_only_code_markdown.code_strings.append(read_only_code_with_imports)
read_only_code_markdown_tokens = len(tokenizer.encode(read_only_code_markdown.markdown))
total_tokens = final_read_writable_tokens + read_only_code_markdown_tokens
if total_tokens <= token_limit:
return CodeString(code=final_read_writable_code).code, read_only_code_markdown.markdown
if read_only_code_with_imports.code:
read_only_code_markdown.code_strings.append(read_only_code_with_imports)
logger.debug("Code context has exceeded token limit, removing read-only code")
return CodeString(code=final_read_writable_code).code, ""
# Extract code from file paths containing helpers of helpers
for file_path, qualified_helper_function_names in helpers_of_helpers_no_overlap.items():
try:
original_code = file_path.read_text("utf8")
except Exception as e:
logger.exception(f"Error while parsing {file_path}: {e}")
continue
try:
read_only_code = get_read_only_code(
original_code, set(), qualified_helper_function_names, remove_docstrings
)
except ValueError as e:
logger.debug(f"Error while getting read-only code: {e}")
continue
read_only_code_with_imports = CodeString(
code=add_needed_imports_from_module(
src_module_code=original_code,
dst_module_code=read_only_code,
src_path=file_path,
dst_path=file_path,
project_root=project_root_path,
helper_functions_fqn=helpers_of_helpers_no_overlap_fqn[file_path],
),
file_path=Path(file_path),
)
if read_only_code_with_imports.code:
read_only_code_markdown.code_strings.append(read_only_code_with_imports)
return read_only_code_markdown
def get_file_path_to_helper_functions_dict(
file_path_to_qualified_function_names: dict[Path, set[str]], project_root_path: Path
) -> tuple[dict[Path, set[str]], dict[Path, set[str]]]:
file_path_to_helper_function_qualified_names = defaultdict(set)
file_path_to_helper_function_fqn = defaultdict(set)
for file_path in file_path_to_qualified_function_names:
script = jedi.Script(path=file_path, project=jedi.Project(path=project_root_path))
file_refs = script.get_names(all_scopes=True, definitions=False, references=True)
for qualified_function_name in file_path_to_qualified_function_names[file_path]:
names = [
ref
for ref in file_refs
if ref.full_name and belongs_to_function_qualified(ref, qualified_function_name)
]
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 definition.type == "function"
and not belongs_to_function_qualified(definition, qualified_function_name)
):
file_path_to_helper_function_qualified_names[definition_path].add(
get_qualified_name(definition.module_name, definition.full_name)
)
file_path_to_helper_function_fqn[definition_path].add(definition.full_name)
return file_path_to_helper_function_qualified_names, file_path_to_helper_function_fqn
def is_dunder_method(name: str) -> bool:
@ -166,7 +245,6 @@ def get_section_names(node: cst.CSTNode) -> list[str]:
def remove_docstring_from_body(indented_block: cst.IndentedBlock) -> cst.CSTNode:
"""Removes the docstring from an indented block if it exists"""
print(indented_block)
if not isinstance(indented_block.body[0], cst.SimpleStatementLine):
return indented_block
first_stmt = indented_block.body[0].body[0]
@ -268,7 +346,11 @@ def get_read_writable_code(code: str, target_functions: set[str]) -> str:
def prune_cst_for_read_only_code(
node: cst.CSTNode, target_functions: set[str], prefix: str = "", remove_docstrings: bool = False
node: cst.CSTNode,
target_functions: set[str],
helpers_of_helper_functions: set[str],
prefix: str = "",
remove_docstrings: bool = False,
) -> tuple[cst.CSTNode | None, bool]:
"""Recursively filter the node for read-only context:
@ -284,6 +366,8 @@ def prune_cst_for_read_only_code(
if isinstance(node, cst.FunctionDef):
qualified_name = f"{prefix}.{node.name.value}" if prefix else node.name.value
# If it's a target function, remove it but mark found_target = True
if qualified_name in helpers_of_helper_functions:
return node, True
if qualified_name in target_functions:
return None, True
# Keep only dunder methods
@ -309,15 +393,9 @@ def prune_cst_for_read_only_code(
new_class_body: list[CSTNode] = []
for stmt in node.body.body:
filtered, found_target = prune_cst_for_read_only_code(
stmt, target_functions, class_prefix, remove_docstrings=remove_docstrings
stmt, target_functions, helpers_of_helper_functions, class_prefix, remove_docstrings=remove_docstrings
)
found_in_class |= found_target
if isinstance(filtered, cst.FunctionDef):
# Check if it's a target or non-dunder method
qname = f"{class_prefix}.{filtered.name.value}"
if qname in target_functions or not is_dunder_method(filtered.name.value):
continue
if filtered:
new_class_body.append(filtered)
@ -345,7 +423,7 @@ def prune_cst_for_read_only_code(
section_found_target = False
for child in original_content:
filtered, found_target = prune_cst_for_read_only_code(
child, target_functions, prefix, remove_docstrings=remove_docstrings
child, target_functions, helpers_of_helper_functions, prefix, remove_docstrings=remove_docstrings
)
if filtered:
new_children.append(filtered)
@ -356,25 +434,30 @@ def prune_cst_for_read_only_code(
updates[section] = new_children
elif original_content is not None:
filtered, found_target = prune_cst_for_read_only_code(
original_content, target_functions, prefix, remove_docstrings=remove_docstrings
original_content,
target_functions,
helpers_of_helper_functions,
prefix,
remove_docstrings=remove_docstrings,
)
found_any_target |= found_target
if filtered:
updates[section] = filtered
if updates:
return (node.with_changes(**updates), found_any_target)
return node, found_any_target
return None, False
def get_read_only_code(code: str, target_functions: set[str], remove_docstrings: bool = False) -> str:
def get_read_only_code(
code: str, target_functions: set[str], helpers_of_helper_functions: set[str], remove_docstrings: bool = False
) -> str:
"""Creates a read-only version of the code by parsing and filtering the code to keep only
class contextual information, and other module scoped variables.
"""
module = cst.parse_module(code)
filtered_node, found_target = prune_cst_for_read_only_code(
module, target_functions, remove_docstrings=remove_docstrings
module, target_functions, helpers_of_helper_functions, remove_docstrings=remove_docstrings
)
if not found_target:
raise ValueError("No target functions found in the provided code")

25
codeflash/optimization/function_context.py
View file

@ -12,7 +12,11 @@ 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 module_name_from_file_path, path_belongs_to_site_packages
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
@ -26,7 +30,7 @@ def belongs_to_method(name: Name, class_name: str, method_name: str) -> bool:
def belongs_to_function(name: Name, function_name: str) -> bool:
"""Check if the given jedi Name is a direct child of the specified function"""
"""Check if the given jedi Name is a direct child of the specified function."""
if name.name == function_name: # Handles function definition and recursive function calls
return False
if name := name.parent():
@ -36,13 +40,28 @@ def belongs_to_function(name: Name, function_name: str) -> bool:
def belongs_to_class(name: Name, class_name: str) -> bool:
"""Check if given jedi Name is a direct child of the specified class"""
"""Check if given jedi Name is a direct child of the specified class."""
while name := name.parent():
if name.type == "class":
return name.name == class_name
return False
def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> bool:
"""Check if the given jedi Name is a direct child of the specified function, matched by qualified function name."""
try:
if get_qualified_name(name.module_name, name.full_name) == qualified_function_name:
# Handles function definition and recursive function calls
return False
if name := name.parent():
if name.type == "function":
return get_qualified_name(name.module_name, name.full_name) == qualified_function_name
return False
except ValueError as e:
logger.exception(f"Error while checking if {name.full_name} belongs to {qualified_function_name}: {e}")
return False
def get_type_annotation_context(
function: FunctionToOptimize, jedi_script: jedi.Script, project_root_path: Path
) -> tuple[list[FunctionSource], set[tuple[str, str]]]:

579
tests/test_code_context_extractor.py
View file

@ -24,6 +24,10 @@ class HelperClass:
return self.name
def main_method():
return "hello"
class MainClass:
def __init__(self, name):
self.name = name
@ -67,7 +71,7 @@ def test_code_replacement10() -> None:
func_top_optimize = FunctionToOptimize(
function_name="main_method", file_path=file_path, parents=[FunctionParent("MainClass", "ClassDef")]
)
original_code = file_path.read_text()
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize=func_top_optimize, project_root_path=file_path.parent
)
@ -90,7 +94,6 @@ def test_code_replacement10() -> None:
```python:{file_path}
from __future__ import annotations
class HelperClass:
def __init__(self, name):
self.name = name
@ -151,7 +154,6 @@ class Graph:
from __future__ import annotations
from collections import defaultdict
class Graph:
def __init__(self, vertices):
self.graph = defaultdict(list)
@ -184,14 +186,8 @@ def test_bubble_sort_helper() -> None:
)
expected_read_write_context = """
from bubble_sort_with_math import sorter
import math
def sort_from_another_file(arr):
sorted_arr = sorter(arr)
return sorted_arr
from bubble_sort_with_math import sorter
def sorter(arr):
arr.sort()
@ -199,6 +195,12 @@ def sorter(arr):
print(x)
return arr
def sort_from_another_file(arr):
sorted_arr = sorter(arr)
return sorted_arr
"""
expected_read_only_context = ""
@ -206,84 +208,6 @@ def sorter(arr):
assert read_only_context.strip() == dedent(expected_read_only_context).strip()
def test_repo_helper() -> None:
path_to_file = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "main.py"
)
path_to_utils = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "utils.py"
)
function_to_optimize = FunctionToOptimize(
function_name="fetch_and_process_data",
file_path=str(path_to_file),
parents=[],
starting_line=None,
ending_line=None,
)
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize, Path(__file__).resolve().parent.parent
)
expected_read_write_context = """
import requests
from globals import API_URL
from utils import DataProcessor
def fetch_and_process_data():
# Use the global variable for the request
response = requests.get(API_URL)
response.raise_for_status()
raw_data = response.text
# Use code from another file (utils.py)
processor = DataProcessor()
processed = processor.process_data(raw_data)
processed = processor.add_prefix(processed)
return processed
class DataProcessor:
def process_data(self, raw_data: str) -> str:
\"\"\"Process raw data by converting it to uppercase.\"\"\"
return raw_data.upper()
def add_prefix(self, data: str, prefix: str = "PREFIX_") -> str:
\"\"\"Add a prefix to the processed data.\"\"\"
return prefix + data
"""
expected_read_only_context = f"""
```python:{path_to_file}
if __name__ == "__main__":
result = fetch_and_process_data()
print("Processed data:", result)
```
```python:{path_to_utils}
import math
class DataProcessor:
\"\"\"A class for processing data.\"\"\"
number = 1
def __init__(self, default_prefix: str = "PREFIX_"):
\"\"\"Initialize the DataProcessor with a default prefix.\"\"\"
self.default_prefix = default_prefix
self.number += math.log(self.number)
def __repr__(self) -> str:
\"\"\"Return a string representation of the DataProcessor.\"\"\"
return f"DataProcessor(default_prefix={{self.default_prefix!r}})"
```
"""
assert read_write_context.strip() == dedent(expected_read_write_context).strip()
assert read_only_context.strip() == dedent(expected_read_only_context).strip()
def test_flavio_typed_code_helper() -> None:
code = '''
@ -569,7 +493,27 @@ class AbstractCacheBackend(CacheBackend, Protocol[_KEY_T, _STORE_T]):
"""Interface for cache backends used by the persistent cache decorator."""
def __init__(self) -> None: ...
def hash_key(
self,
*,
func: Callable[_P, Any],
args: tuple[Any, ...],
kwargs: dict[str, Any],
) -> tuple[str, _KEY_T]: ...
def encode(self, *, data: Any) -> _STORE_T: # noqa: ANN401
...
def decode(self, *, data: _STORE_T) -> Any: # noqa: ANN401
...
def get(self, *, key: tuple[str, _KEY_T]) -> tuple[datetime.datetime, _STORE_T] | None: ...
def delete(self, *, key: tuple[str, _KEY_T]) -> None: ...
def put(self, *, key: tuple[str, _KEY_T], data: _STORE_T) -> None: ...
_P = ParamSpec("_P")
_R = TypeVar("_R")
_CacheBackendT = TypeVar("_CacheBackendT", bound=CacheBackend)
@ -885,3 +829,462 @@ class HelperClass:
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize, opt.args.project_root
)
def test_repo_helper() -> None:
path_to_file = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "main.py"
)
path_to_utils = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "utils.py"
)
function_to_optimize = FunctionToOptimize(
function_name="fetch_and_process_data",
file_path=str(path_to_file),
parents=[],
starting_line=None,
ending_line=None,
)
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize,
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever",
)
expected_read_write_context = """
import requests
from globals import API_URL
from utils import DataProcessor
class DataProcessor:
def process_data(self, raw_data: str) -> str:
\"\"\"Process raw data by converting it to uppercase.\"\"\"
return raw_data.upper()
def add_prefix(self, data: str, prefix: str = "PREFIX_") -> str:
\"\"\"Add a prefix to the processed data.\"\"\"
return prefix + data
def fetch_and_process_data():
# Use the global variable for the request
response = requests.get(API_URL)
response.raise_for_status()
raw_data = response.text
# Use code from another file (utils.py)
processor = DataProcessor()
processed = processor.process_data(raw_data)
processed = processor.add_prefix(processed)
return processed
"""
expected_read_only_context = f"""
```python:{path_to_utils}
import math
GLOBAL_VAR = 10
class DataProcessor:
\"\"\"A class for processing data.\"\"\"
number = 1
def __init__(self, default_prefix: str = "PREFIX_"):
\"\"\"Initialize the DataProcessor with a default prefix.\"\"\"
self.default_prefix = default_prefix
self.number += math.log(self.number)
def __repr__(self) -> str:
\"\"\"Return a string representation of the DataProcessor.\"\"\"
return f"DataProcessor(default_prefix={{self.default_prefix!r}})"
```
```python:{path_to_file}
if __name__ == "__main__":
result = fetch_and_process_data()
print("Processed data:", result)
```
"""
assert read_write_context.strip() == dedent(expected_read_write_context).strip()
assert read_only_context.strip() == dedent(expected_read_only_context).strip()
def test_repo_helper_of_helper() -> None:
path_to_file = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "main.py"
)
path_to_utils = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "utils.py"
)
path_to_transform_utils = (
Path(__file__).resolve().parent.parent
/ "code_to_optimize"
/ "code_directories"
/ "retriever"
/ "transform_utils.py"
)
function_to_optimize = FunctionToOptimize(
function_name="fetch_and_transform_data",
file_path=str(path_to_file),
parents=[],
starting_line=None,
ending_line=None,
)
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize,
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever",
)
expected_read_write_context = """
from transform_utils import DataTransformer
import requests
from globals import API_URL
from utils import DataProcessor
class DataProcessor:
def process_data(self, raw_data: str) -> str:
\"\"\"Process raw data by converting it to uppercase.\"\"\"
return raw_data.upper()
def transform_data(self, data: str) -> str:
\"\"\"Transform the processed data\"\"\"
return DataTransformer().transform(data)
def fetch_and_transform_data():
# Use the global variable for the request
response = requests.get(API_URL)
raw_data = response.text
# Use code from another file (utils.py)
processor = DataProcessor()
processed = processor.process_data(raw_data)
transformed = processor.transform_data(processed)
return transformed
"""
expected_read_only_context = f"""
```python:{path_to_utils}
import math
GLOBAL_VAR = 10
class DataProcessor:
\"\"\"A class for processing data.\"\"\"
number = 1
def __init__(self, default_prefix: str = "PREFIX_"):
\"\"\"Initialize the DataProcessor with a default prefix.\"\"\"
self.default_prefix = default_prefix
self.number += math.log(self.number)
def __repr__(self) -> str:
\"\"\"Return a string representation of the DataProcessor.\"\"\"
return f"DataProcessor(default_prefix={{self.default_prefix!r}})"
```
```python:{path_to_file}
if __name__ == "__main__":
result = fetch_and_process_data()
print("Processed data:", result)
```
```python:{path_to_transform_utils}
class DataTransformer:
def __init__(self):
self.data = None
def transform(self, data):
self.data = data
return self.data
```
"""
assert read_write_context.strip() == dedent(expected_read_write_context).strip()
assert read_only_context.strip() == dedent(expected_read_only_context).strip()
def test_repo_helper_of_helper_same_class() -> None:
path_to_utils = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "utils.py"
)
path_to_transform_utils = (
Path(__file__).resolve().parent.parent
/ "code_to_optimize"
/ "code_directories"
/ "retriever"
/ "transform_utils.py"
)
function_to_optimize = FunctionToOptimize(
function_name="transform_data_own_method",
file_path=str(path_to_utils),
parents=[FunctionParent(name="DataProcessor", type="ClassDef")],
starting_line=None,
ending_line=None,
)
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize,
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever",
)
expected_read_write_context = """
from transform_utils import DataTransformer
class DataTransformer:
def transform_using_own_method(self, data):
return self.transform(data)
class DataProcessor:
def transform_data_own_method(self, data: str) -> str:
\"\"\"Transform the processed data using own method\"\"\"
return DataTransformer().transform_using_own_method(data)
"""
expected_read_only_context = f"""
```python:{path_to_transform_utils}
class DataTransformer:
def __init__(self):
self.data = None
def transform(self, data):
self.data = data
return self.data
```
```python:{path_to_utils}
import math
GLOBAL_VAR = 10
class DataProcessor:
\"\"\"A class for processing data.\"\"\"
number = 1
def __init__(self, default_prefix: str = "PREFIX_"):
\"\"\"Initialize the DataProcessor with a default prefix.\"\"\"
self.default_prefix = default_prefix
self.number += math.log(self.number)
def __repr__(self) -> str:
\"\"\"Return a string representation of the DataProcessor.\"\"\"
return f"DataProcessor(default_prefix={{self.default_prefix!r}})"
```
"""
assert read_write_context.strip() == dedent(expected_read_write_context).strip()
assert read_only_context.strip() == dedent(expected_read_only_context).strip()
def test_repo_helper_of_helper_same_file() -> None:
path_to_utils = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "utils.py"
)
path_to_transform_utils = (
Path(__file__).resolve().parent.parent
/ "code_to_optimize"
/ "code_directories"
/ "retriever"
/ "transform_utils.py"
)
function_to_optimize = FunctionToOptimize(
function_name="transform_data_same_file_function",
file_path=str(path_to_utils),
parents=[FunctionParent(name="DataProcessor", type="ClassDef")],
starting_line=None,
ending_line=None,
)
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize,
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever",
)
expected_read_write_context = """
from transform_utils import DataTransformer
class DataTransformer:
def transform_using_same_file_function(self, data):
return update_data(data)
class DataProcessor:
def transform_data_same_file_function(self, data: str) -> str:
\"\"\"Transform the processed data using a function from the same file\"\"\"
return DataTransformer().transform_using_same_file_function(data)
"""
expected_read_only_context = f"""
```python:{path_to_transform_utils}
class DataTransformer:
def __init__(self):
self.data = None
def update_data(data):
return data + " updated"
```
```python:{path_to_utils}
import math
GLOBAL_VAR = 10
class DataProcessor:
\"\"\"A class for processing data.\"\"\"
number = 1
def __init__(self, default_prefix: str = "PREFIX_"):
\"\"\"Initialize the DataProcessor with a default prefix.\"\"\"
self.default_prefix = default_prefix
self.number += math.log(self.number)
def __repr__(self) -> str:
\"\"\"Return a string representation of the DataProcessor.\"\"\"
return f"DataProcessor(default_prefix={{self.default_prefix!r}})"
```
"""
assert read_write_context.strip() == dedent(expected_read_write_context).strip()
assert read_only_context.strip() == dedent(expected_read_only_context).strip()
def test_repo_helper_all_same_file() -> None:
path_to_transform_utils = (
Path(__file__).resolve().parent.parent
/ "code_to_optimize"
/ "code_directories"
/ "retriever"
/ "transform_utils.py"
)
function_to_optimize = FunctionToOptimize(
function_name="transform_data_all_same_file",
file_path=str(path_to_transform_utils),
parents=[FunctionParent(name="DataTransformer", type="ClassDef")],
starting_line=None,
ending_line=None,
)
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize,
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever",
)
expected_read_write_context = """
class DataTransformer:
def transform_using_own_method(self, data):
return self.transform(data)
def transform_data_all_same_file(self, data):
new_data = update_data(data)
return self.transform_using_own_method(new_data)
def update_data(data):
return data + " updated"
"""
expected_read_only_context = f"""
```python:{path_to_transform_utils}
class DataTransformer:
def __init__(self):
self.data = None
def transform(self, data):
self.data = data
return self.data
```
"""
assert read_write_context.strip() == dedent(expected_read_write_context).strip()
assert read_only_context.strip() == dedent(expected_read_only_context).strip()
def test_repo_helper_circular_dependency() -> None:
path_to_transform_utils = (
Path(__file__).resolve().parent.parent
/ "code_to_optimize"
/ "code_directories"
/ "retriever"
/ "transform_utils.py"
)
path_to_utils = (
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever" / "utils.py"
)
function_to_optimize = FunctionToOptimize(
function_name="circular_dependency",
file_path=str(path_to_transform_utils),
parents=[FunctionParent(name="DataTransformer", type="ClassDef")],
starting_line=None,
ending_line=None,
)
read_write_context, read_only_context = get_code_optimization_context(
function_to_optimize,
Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "retriever",
)
expected_read_write_context = """
from transform_utils import DataTransformer
from code_to_optimize.code_directories.retriever.utils import DataProcessor
class DataProcessor:
def circular_dependency(self, data: str) -> str:
\"\"\"Test circular dependency\"\"\"
return DataTransformer().circular_dependency(data)
class DataTransformer:
def circular_dependency(self, data):
return DataProcessor().circular_dependency(data)
"""
expected_read_only_context = f"""
```python:{path_to_utils}
import math
GLOBAL_VAR = 10
class DataProcessor:
\"\"\"A class for processing data.\"\"\"
number = 1
def __init__(self, default_prefix: str = "PREFIX_"):
\"\"\"Initialize the DataProcessor with a default prefix.\"\"\"
self.default_prefix = default_prefix
self.number += math.log(self.number)
def __repr__(self) -> str:
\"\"\"Return a string representation of the DataProcessor.\"\"\"
return f"DataProcessor(default_prefix={{self.default_prefix!r}})"
```
```python:{path_to_transform_utils}
class DataTransformer:
def __init__(self):
self.data = None
```
"""
assert read_write_context.strip() == dedent(expected_read_write_context).strip()
assert read_only_context.strip() == dedent(expected_read_only_context).strip()

44
tests/test_get_read_only_code.py
View file

@ -21,7 +21,7 @@ def test_basic_class() -> None:
class_var = "value"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -47,7 +47,7 @@ def test_dunder_methods() -> None:
return f"Value: {self.x}"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -75,7 +75,7 @@ def test_dunder_methods_remove_docstring() -> None:
return f"Value: {self.x}"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, remove_docstrings=True)
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set(), remove_docstrings=True)
assert dedent(expected).strip() == output.strip()
@ -102,7 +102,7 @@ def test_class_remove_docstring() -> None:
return f"Value: {self.x}"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, remove_docstrings=True)
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set(), remove_docstrings=True)
assert dedent(expected).strip() == output.strip()
@ -131,7 +131,7 @@ def test_mixed_remove_docstring() -> None:
return f"Value: {self.x}"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, remove_docstrings=True)
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set(), remove_docstrings=True)
assert dedent(expected).strip() == output.strip()
@ -149,7 +149,7 @@ def test_target_in_nested_class() -> None:
"""
with pytest.raises(ValueError, match="No target functions found in the provided code"):
get_read_only_code(dedent(code), {"Outer.Inner.target_method"})
get_read_only_code(dedent(code), {"Outer.Inner.target_method"}, set())
def test_docstrings() -> None:
@ -171,7 +171,7 @@ def test_docstrings() -> None:
\"\"\"Class docstring.\"\"\"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -190,7 +190,7 @@ def test_method_signatures() -> None:
expected = """"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -214,7 +214,7 @@ def test_multiple_top_level_targets() -> None:
self.x = 42
"""
output = get_read_only_code(dedent(code), {"TestClass.target1", "TestClass.target2"})
output = get_read_only_code(dedent(code), {"TestClass.target1", "TestClass.target2"}, set())
assert dedent(expected).strip() == output.strip()
@ -234,7 +234,7 @@ def test_class_annotations() -> None:
var2: str
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -256,7 +256,7 @@ def test_class_annotations_if() -> None:
var2: str
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -282,7 +282,7 @@ def test_class_annotations_try() -> None:
continue
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -318,7 +318,7 @@ def test_class_annotations_else() -> None:
var2: str
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -333,7 +333,7 @@ def test_top_level_functions() -> None:
expected = """"""
output = get_read_only_code(dedent(code), {"target_function"})
output = get_read_only_code(dedent(code), {"target_function"}, set())
assert dedent(expected).strip() == output.strip()
@ -352,7 +352,7 @@ def test_module_var() -> None:
x = 5
"""
output = get_read_only_code(dedent(code), {"target_function"})
output = get_read_only_code(dedent(code), {"target_function"}, set())
assert dedent(expected).strip() == output.strip()
@ -379,7 +379,7 @@ def test_module_var_if() -> None:
z = 10
"""
output = get_read_only_code(dedent(code), {"target_function"})
output = get_read_only_code(dedent(code), {"target_function"}, set())
assert dedent(expected).strip() == output.strip()
@ -414,7 +414,7 @@ def test_conditional_class_definitions() -> None:
platform = "other"
"""
output = get_read_only_code(dedent(code), {"PlatformClass.target_method"})
output = get_read_only_code(dedent(code), {"PlatformClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -473,7 +473,7 @@ def test_multiple_except_clauses() -> None:
error_type = "cleanup"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -526,7 +526,7 @@ def test_with_statement_and_loops() -> None:
context = "cleanup"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -575,7 +575,7 @@ def test_async_with_try_except() -> None:
status = "cancelled"
"""
output = get_read_only_code(dedent(code), {"TestClass.target_method"})
output = get_read_only_code(dedent(code), {"TestClass.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -685,7 +685,7 @@ def test_simplified_complete_implementation() -> None:
self.error = str(e)
"""
output = get_read_only_code(dedent(code), {"DataProcessor.target_method", "ResultHandler.target_method"})
output = get_read_only_code(dedent(code), {"DataProcessor.target_method", "ResultHandler.target_method"}, set())
assert dedent(expected).strip() == output.strip()
@ -795,6 +795,6 @@ def test_simplified_complete_implementation_no_docstring() -> None:
"""
output = get_read_only_code(
dedent(code), {"DataProcessor.target_method", "ResultHandler.target_method"}, remove_docstrings=True
dedent(code), {"DataProcessor.target_method", "ResultHandler.target_method"}, set(), remove_docstrings=True
)
assert dedent(expected).strip() == output.strip()