Skip to content

Commit 4396abc

Browse files
committed
cleaning up
1 parent f3cba99 commit 4396abc

File tree

2 files changed

+15
-163
lines changed

2 files changed

+15
-163
lines changed

codeflash/code_utils/edit_generated_tests.py

Lines changed: 14 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,17 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
3838
self.context_stack.pop()
3939
return node
4040

41+
def get_comment(self, match_key: str) -> str:
42+
# calculate speedup and output comment
43+
original_time = self.original_runtimes[match_key]
44+
optimized_time = self.optimized_runtimes[match_key]
45+
perf_gain = format_perf(
46+
abs(performance_gain(original_runtime_ns=original_time, optimized_runtime_ns=optimized_time) * 100)
47+
)
48+
status = "slower" if optimized_time > original_time else "faster"
49+
# Create the runtime comment
50+
return f"# {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})"
51+
4152
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
4253
self.context_stack.append(node.name)
4354
i = len(node.body) - 1
@@ -55,41 +66,13 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
5566
inv_id = str(i) + "_" + str(j)
5667
match_key = key + "#" + inv_id
5768
if match_key in self.original_runtimes and match_key in self.optimized_runtimes:
58-
# calculate speedup and output comment
59-
original_time = self.original_runtimes[match_key]
60-
optimized_time = self.optimized_runtimes[match_key]
61-
perf_gain = format_perf(
62-
abs(
63-
performance_gain(
64-
original_runtime_ns=original_time, optimized_runtime_ns=optimized_time
65-
)
66-
* 100
67-
)
68-
)
69-
status = "slower" if optimized_time > original_time else "faster"
70-
# Create the runtime comment
71-
comment_text = f"# {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})"
72-
self.results[internal_node.lineno] = comment_text
69+
self.results[internal_node.lineno] = self.get_comment(match_key)
7370
j -= 1
7471
else:
7572
inv_id = str(i)
7673
match_key = key + "#" + inv_id
7774
if match_key in self.original_runtimes and match_key in self.optimized_runtimes:
78-
# calculate speedup and output comment
79-
original_time = self.original_runtimes[match_key]
80-
optimized_time = self.optimized_runtimes[match_key]
81-
perf_gain = format_perf(
82-
abs(
83-
performance_gain(original_runtime_ns=original_time, optimized_runtime_ns=optimized_time)
84-
* 100
85-
)
86-
)
87-
status = "slower" if optimized_time > original_time else "faster"
88-
# Create the runtime comment
89-
comment_text = (
90-
f"# {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})"
91-
)
92-
self.results[line_node.lineno] = comment_text
75+
self.results[line_node.lineno] = self.get_comment(match_key)
9376
i -= 1
9477
self.context_stack.pop()
9578
return node
@@ -106,7 +89,7 @@ def get_fn_call_linenos(
10689

10790

10891
class CommentAdder(cst.CSTTransformer):
109-
"""Transformer that adds comment 'a' to specified lines."""
92+
"""Transformer that adds comments to specified lines."""
11093

11194
# Declare metadata dependencies
11295
METADATA_DEPENDENCIES = (PositionProvider,)

tests/test_add_runtime_comments.py

Lines changed: 1 addition & 132 deletions
Original file line numberDiff line numberDiff line change
@@ -1902,135 +1902,4 @@ def test_bubble_sort(input, expected_output):
19021902

19031903
# Check that comments were added
19041904
modified_source = result.generated_tests[0].generated_original_test_source
1905-
assert modified_source == expected
1906-
1907-
"""TODO Future tests"""
1908-
# def test_runtime_comment_addition_else(self, test_config):
1909-
# """Test basic functionality of adding runtime comments."""
1910-
# # Create test source code
1911-
# os.chdir(test_config.project_root_path)
1912-
# test_source = """def test_bubble_sort():
1913-
# i = 0
1914-
# if 1>2:
1915-
# b = 3
1916-
# else:
1917-
# b1 = 6
1918-
# codeflash_output = bubble_sort([3, 1, 2])
1919-
# assert codeflash_output == [1, 2, 3]
1920-
# i += 1
1921-
# d = 5
1922-
# """
1923-
# expected = """def test_bubble_sort():
1924-
# i = 0
1925-
# if 1>2:
1926-
# b = 3
1927-
# else:
1928-
# b1 = 6
1929-
# codeflash_output = bubble_sort([3, 1, 2]) # 500μs -> 300μs (66.7% faster)
1930-
# assert codeflash_output == [1, 2, 3]
1931-
# i += 1
1932-
# d = 5
1933-
# """
1934-
# generated_test = GeneratedTests(
1935-
# generated_original_test_source=test_source,
1936-
# instrumented_behavior_test_source="",
1937-
# instrumented_perf_test_source="",
1938-
# behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
1939-
# perf_file_path=test_config.tests_root / "test_perf.py",
1940-
# )
1941-
# generated_tests = GeneratedTestsList(generated_tests=[generated_test])
1942-
#
1943-
# # Create test results
1944-
# original_test_results = TestResults()
1945-
# optimized_test_results = TestResults()
1946-
#
1947-
# # Add test invocations with different runtimes
1948-
# original_invocation1 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='2_1_0') # 500μs
1949-
# optimized_invocation1 = self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='2_1_0') # 300μs
1950-
# # longer runtime than minimum, will not contribute
1951-
# original_invocation2 = self.create_test_invocation("test_bubble_sort", 600_000, iteration_id='2_1_1') # 500μs
1952-
# optimized_invocation2 = self.create_test_invocation("test_bubble_sort", 400_000, iteration_id='2_1_1') # 300μs
1953-
# original_invocation3 = self.create_test_invocation("test_bubble_sort", 700_000, iteration_id='2_1_2') # 500μs
1954-
# optimized_invocation3 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='2_1_2') # 300μs
1955-
#
1956-
# original_test_results.add(original_invocation1)
1957-
# optimized_test_results.add(optimized_invocation1)
1958-
# original_test_results.add(original_invocation2)
1959-
# optimized_test_results.add(optimized_invocation2)
1960-
# original_test_results.add(original_invocation3)
1961-
# optimized_test_results.add(optimized_invocation3)
1962-
# original_runtimes = original_test_results.usable_runtime_data_by_test_case()
1963-
# optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case()
1964-
# # Test the functionality
1965-
# result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
1966-
#
1967-
# # Check that comments were added
1968-
# modified_source = result.generated_tests[0].generated_original_test_source
1969-
# assert modified_source == expected
1970-
#
1971-
# def test_runtime_comment_addition_elif(self, test_config):
1972-
# """Test basic functionality of adding runtime comments."""
1973-
# # Create test source code
1974-
# os.chdir(test_config.project_root_path)
1975-
# test_source = """def test_bubble_sort():
1976-
# i = 0
1977-
# if 1>2:
1978-
# b = 3
1979-
# elif 2<3:
1980-
# b1 = 6
1981-
# codeflash_output = bubble_sort([3, 1, 2])
1982-
# assert codeflash_output == [1, 2, 3]
1983-
# i += 1
1984-
# else:
1985-
# qwe = 1
1986-
# d = 5
1987-
# """
1988-
# expected = """def test_bubble_sort():
1989-
# i = 0
1990-
# if 1>2:
1991-
# b = 3
1992-
# elif 2<3:
1993-
# b1 = 6
1994-
# codeflash_output = bubble_sort([3, 1, 2]) # 500μs -> 300μs (66.7% faster)
1995-
# assert codeflash_output == [1, 2, 3]
1996-
# i += 1
1997-
# else:
1998-
# qwe = 1
1999-
# d = 5
2000-
# """
2001-
# generated_test = GeneratedTests(
2002-
# generated_original_test_source=test_source,
2003-
# instrumented_behavior_test_source="",
2004-
# instrumented_perf_test_source="",
2005-
# behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
2006-
# perf_file_path=test_config.tests_root / "test_perf.py",
2007-
# )
2008-
# generated_tests = GeneratedTestsList(generated_tests=[generated_test])
2009-
#
2010-
# # Create test results
2011-
# original_test_results = TestResults()
2012-
# optimized_test_results = TestResults()
2013-
#
2014-
# # Add test invocations with different runtimes
2015-
# original_invocation1 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='2_1_0') # 500μs
2016-
# optimized_invocation1 = self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='2_1_0') # 300μs
2017-
# # longer runtime than minimum, will not contribute
2018-
# original_invocation2 = self.create_test_invocation("test_bubble_sort", 600_000, iteration_id='2_1_1') # 500μs
2019-
# optimized_invocation2 = self.create_test_invocation("test_bubble_sort", 400_000, iteration_id='2_1_1') # 300μs
2020-
# original_invocation3 = self.create_test_invocation("test_bubble_sort", 700_000, iteration_id='2_1_2') # 500μs
2021-
# optimized_invocation3 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='2_1_2') # 300μs
2022-
#
2023-
# original_test_results.add(original_invocation1)
2024-
# optimized_test_results.add(optimized_invocation1)
2025-
# original_test_results.add(original_invocation2)
2026-
# optimized_test_results.add(optimized_invocation2)
2027-
# original_test_results.add(original_invocation3)
2028-
# optimized_test_results.add(optimized_invocation3)
2029-
# original_runtimes = original_test_results.usable_runtime_data_by_test_case()
2030-
# optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case()
2031-
# # Test the functionality
2032-
# result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
2033-
#
2034-
# # Check that comments were added
2035-
# modified_source = result.generated_tests[0].generated_original_test_source
2036-
# assert modified_source == expected
1905+
assert modified_source == expected

0 commit comments

Comments
 (0)