381 lines
12 KiB
Python
381 lines
12 KiB
Python
import pathlib
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
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
|
|
|
|
|
|
def calculate_something(data):
|
|
return data + 1
|
|
|
|
|
|
def simple_function_with_one_dep(data):
|
|
return calculate_something(data)
|
|
|
|
|
|
def test_simple_dependencies() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
helper_functions = get_function_variables_definitions(
|
|
FunctionToOptimize("simple_function_with_one_dep", str(file_path), []), str(file_path.parent.resolve())
|
|
)[0]
|
|
assert len(helper_functions) == 1
|
|
assert helper_functions[0].jedi_definition.full_name == "test_function_dependencies.calculate_something"
|
|
|
|
|
|
def global_dependency_1(num):
|
|
return num + 1
|
|
|
|
|
|
def global_dependency_2(num):
|
|
return num + 1
|
|
|
|
|
|
def global_dependency_3(num):
|
|
return num + 1
|
|
|
|
|
|
class A:
|
|
def calculate_something_1(self, num):
|
|
return num + 1
|
|
|
|
def run(self):
|
|
a = 1
|
|
b = self.calculate_something_1(a)
|
|
c = global_dependency_1(b)
|
|
return c
|
|
|
|
def function_in_list_comprehension(self):
|
|
return [global_dependency_3(1) for x in range(10)]
|
|
|
|
def add_two(self, num):
|
|
return num + 2
|
|
|
|
def method_in_list_comprehension(self):
|
|
return [self.add_two(1) for x in range(10)]
|
|
|
|
def nested_function(self):
|
|
def nested():
|
|
return global_dependency_3(1)
|
|
|
|
return nested() + self.add_two(3)
|
|
|
|
|
|
class B:
|
|
def calculate_something_2(self, num):
|
|
return num + 1
|
|
|
|
def run(self):
|
|
a = 1
|
|
b = self.calculate_something_2(a)
|
|
c = global_dependency_2(b)
|
|
return c
|
|
|
|
|
|
class C:
|
|
def calculate_something_3(self, num):
|
|
return num + 1
|
|
|
|
def run(self):
|
|
a = 1
|
|
b = self.calculate_something_3(a)
|
|
c = global_dependency_3(b)
|
|
return c
|
|
|
|
def recursive(self, num):
|
|
if num == 0:
|
|
return 0
|
|
num_1 = self.calculate_something_3(num)
|
|
return self.recursive(num) + num_1
|
|
|
|
|
|
def test_multiple_classes_dependencies() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
helper_functions = get_function_variables_definitions(
|
|
FunctionToOptimize("run", str(file_path), [FunctionParent("C", "ClassDef")]), str(file_path.parent.resolve())
|
|
)
|
|
|
|
assert len(helper_functions) == 2
|
|
assert list(map(lambda x: x.fully_qualified_name, helper_functions[0])) == [
|
|
"test_function_dependencies.global_dependency_3",
|
|
"test_function_dependencies.C.calculate_something_3",
|
|
]
|
|
|
|
|
|
def recursive_dependency_1(num):
|
|
if num == 0:
|
|
return 0
|
|
num_1 = calculate_something(num)
|
|
return recursive_dependency_1(num) + num_1
|
|
|
|
|
|
def test_recursive_dependency() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
helper_functions = get_function_variables_definitions(
|
|
FunctionToOptimize("recursive_dependency_1", str(file_path), []), str(file_path.parent.resolve())
|
|
)[0]
|
|
assert len(helper_functions) == 1
|
|
assert helper_functions[0].jedi_definition.full_name == "test_function_dependencies.calculate_something"
|
|
assert helper_functions[0].fully_qualified_name == "test_function_dependencies.calculate_something"
|
|
|
|
|
|
@dataclass
|
|
class MyData:
|
|
MyInt: int
|
|
|
|
|
|
def calculate_something_ann(data):
|
|
return data + 1
|
|
|
|
|
|
def simple_function_with_one_dep_ann(data: MyData):
|
|
return calculate_something_ann(data)
|
|
|
|
|
|
def list_comprehension_dependency(data: MyData):
|
|
return [calculate_something(data) for x in range(10)]
|
|
|
|
|
|
def test_simple_dependencies_ann() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
helper_functions = get_function_variables_definitions(
|
|
FunctionToOptimize("simple_function_with_one_dep_ann", str(file_path), []), str(file_path.parent.resolve())
|
|
)[0]
|
|
assert len(helper_functions) == 2
|
|
assert helper_functions[0].jedi_definition.full_name == "test_function_dependencies.MyData"
|
|
assert helper_functions[1].jedi_definition.full_name == "test_function_dependencies.calculate_something_ann"
|
|
|
|
|
|
from collections import defaultdict
|
|
|
|
|
|
class Graph:
|
|
def __init__(self, vertices):
|
|
self.graph = defaultdict(list)
|
|
self.V = vertices # No. of vertices
|
|
|
|
def addEdge(self, u, v):
|
|
self.graph[u].append(v)
|
|
|
|
def topologicalSortUtil(self, v, visited, stack):
|
|
visited[v] = True
|
|
|
|
for i in self.graph[v]:
|
|
if visited[i] == False:
|
|
self.topologicalSortUtil(i, visited, stack)
|
|
|
|
stack.insert(0, v)
|
|
|
|
def topologicalSort(self):
|
|
visited = [False] * self.V
|
|
stack = []
|
|
|
|
for i in range(self.V):
|
|
if visited[i] == False:
|
|
self.topologicalSortUtil(i, visited, stack)
|
|
|
|
# Print contents of stack
|
|
return stack
|
|
|
|
|
|
def test_class_method_dependencies() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
|
|
function_to_optimize = FunctionToOptimize(
|
|
function_name="topologicalSort",
|
|
file_path=str(file_path),
|
|
parents=[FunctionParent(name="Graph", type="ClassDef")],
|
|
starting_line=None,
|
|
ending_line=None,
|
|
)
|
|
func_optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=TestConfig(
|
|
tests_root=file_path,
|
|
tests_project_rootdir=file_path.parent,
|
|
project_root_path=file_path.parent,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
),
|
|
)
|
|
with open(file_path) as f:
|
|
original_code = f.read()
|
|
ctx_result = func_optimizer.get_code_optimization_context()
|
|
if not is_successful(ctx_result):
|
|
pytest.fail()
|
|
code_context = ctx_result.unwrap()
|
|
# 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"
|
|
)
|
|
assert code_context.helper_functions[0].jedi_definition.name == "topologicalSortUtil"
|
|
assert (
|
|
code_context.helper_functions[0].fully_qualified_name == "test_function_dependencies.Graph.topologicalSortUtil"
|
|
)
|
|
assert code_context.helper_functions[0].qualified_name == "Graph.topologicalSortUtil"
|
|
assert (
|
|
code_context.code_to_optimize_with_helpers
|
|
== """from collections import defaultdict
|
|
|
|
class Graph:
|
|
def __init__(self, vertices):
|
|
self.graph = defaultdict(list)
|
|
self.V = vertices # No. of vertices
|
|
def topologicalSortUtil(self, v, visited, stack):
|
|
visited[v] = True
|
|
|
|
for i in self.graph[v]:
|
|
if visited[i] == False:
|
|
self.topologicalSortUtil(i, visited, stack)
|
|
|
|
stack.insert(0, v)
|
|
def topologicalSort(self):
|
|
visited = [False] * self.V
|
|
stack = []
|
|
|
|
for i in range(self.V):
|
|
if visited[i] == False:
|
|
self.topologicalSortUtil(i, visited, stack)
|
|
|
|
# Print contents of stack
|
|
return stack
|
|
"""
|
|
)
|
|
|
|
|
|
def calculate_something_else(data):
|
|
return data + 1
|
|
|
|
|
|
def imalittledecorator(func):
|
|
def wrapper(data):
|
|
return func(data)
|
|
|
|
return wrapper
|
|
|
|
|
|
@imalittledecorator
|
|
def simple_function_with_decorator_dep(data):
|
|
return calculate_something_else(data)
|
|
|
|
|
|
@pytest.mark.skip(reason="no decorator dependency support")
|
|
def test_decorator_dependencies() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
helper_functions = get_function_variables_definitions(
|
|
FunctionToOptimize("simple_function_with_decorator_dep", str(file_path), []), str(file_path.parent.resolve())
|
|
)[0]
|
|
assert len(helper_functions) == 2
|
|
assert {helper_functions[0][0].definition.full_name, helper_functions[1][0].definition.full_name} == {
|
|
"test_function_dependencies.calculate_something",
|
|
"test_function_dependencies.imalittledecorator",
|
|
}
|
|
|
|
|
|
def test_recursive_function_context() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
|
|
function_to_optimize = FunctionToOptimize(
|
|
function_name="recursive",
|
|
file_path=str(file_path),
|
|
parents=[FunctionParent(name="C", type="ClassDef")],
|
|
starting_line=None,
|
|
ending_line=None,
|
|
)
|
|
func_optimizer = FunctionOptimizer(
|
|
function_to_optimize=function_to_optimize,
|
|
test_cfg=TestConfig(
|
|
tests_root=file_path,
|
|
tests_project_rootdir=file_path.parent,
|
|
project_root_path=file_path.parent,
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
),
|
|
)
|
|
with open(file_path) as f:
|
|
original_code = f.read()
|
|
|
|
ctx_result = func_optimizer.get_code_optimization_context()
|
|
if not is_successful(ctx_result):
|
|
pytest.fail()
|
|
code_context = ctx_result.unwrap()
|
|
assert len(code_context.helper_functions) == 2
|
|
assert code_context.helper_functions[0].fully_qualified_name == "test_function_dependencies.C.calculate_something_3"
|
|
assert code_context.helper_functions[1].fully_qualified_name == "test_function_dependencies.C.recursive"
|
|
assert (
|
|
code_context.code_to_optimize_with_helpers
|
|
== """class C:
|
|
def calculate_something_3(self, num):
|
|
return num + 1
|
|
def recursive(self, num):
|
|
if num == 0:
|
|
return 0
|
|
num_1 = self.calculate_something_3(num)
|
|
return self.recursive(num) + num_1
|
|
"""
|
|
)
|
|
|
|
|
|
def test_list_comprehension_dependency() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
helper_functions = get_function_variables_definitions(
|
|
FunctionToOptimize("list_comprehension_dependency", str(file_path), []), str(file_path.parent.resolve())
|
|
)[0]
|
|
assert len(helper_functions) == 2
|
|
assert helper_functions[0].jedi_definition.full_name == "test_function_dependencies.MyData"
|
|
assert helper_functions[1].jedi_definition.full_name == "test_function_dependencies.calculate_something"
|
|
|
|
|
|
def test_function_in_method_list_comprehension() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
function_to_optimize = FunctionToOptimize(
|
|
function_name="function_in_list_comprehension",
|
|
file_path=str(file_path),
|
|
parents=[FunctionParent(name="A", type="ClassDef")],
|
|
starting_line=None,
|
|
ending_line=None,
|
|
)
|
|
|
|
helper_functions = get_function_variables_definitions(function_to_optimize, str(file_path.parent.resolve()))[0]
|
|
|
|
assert len(helper_functions) == 1
|
|
assert helper_functions[0].jedi_definition.full_name == "test_function_dependencies.global_dependency_3"
|
|
|
|
|
|
def test_method_in_method_list_comprehension() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
function_to_optimize = FunctionToOptimize(
|
|
function_name="method_in_list_comprehension",
|
|
file_path=str(file_path),
|
|
parents=[FunctionParent(name="A", type="ClassDef")],
|
|
starting_line=None,
|
|
ending_line=None,
|
|
)
|
|
|
|
helper_functions = get_function_variables_definitions(function_to_optimize, str(file_path.parent.resolve()))[0]
|
|
|
|
assert len(helper_functions) == 1
|
|
assert helper_functions[0].jedi_definition.full_name == "test_function_dependencies.A.add_two"
|
|
|
|
|
|
def test_nested_method() -> None:
|
|
file_path = pathlib.Path(__file__).resolve()
|
|
function_to_optimize = FunctionToOptimize(
|
|
function_name="nested_function",
|
|
file_path=str(file_path),
|
|
parents=[FunctionParent(name="A", type="ClassDef")],
|
|
starting_line=None,
|
|
ending_line=None,
|
|
)
|
|
|
|
helper_functions = get_function_variables_definitions(function_to_optimize, str(file_path.parent.resolve()))[0]
|
|
|
|
# The nested function should be included in the helper functions
|
|
assert len(helper_functions) == 1
|
|
assert helper_functions[0].jedi_definition.full_name == "test_function_dependencies.A.add_two"
|