Reworked testgen-context retrieval

code_to_optimize/bubble_sort2.py

@@ -1,6 +1,3 @@
 def sorter(arr):
     arr.sort()
-    return arr
-
-
-CACHED_TESTS = "import unittest\ndef sorter(arr):\n    for i in range(len(arr)):\n        for j in range(len(arr) - 1):\n            if arr[j] > arr[j + 1]:\n                temp = arr[j]\n                arr[j] = arr[j + 1]\n                arr[j + 1] = temp\n    return arr\nclass SorterTestCase(unittest.TestCase):\n    def test_empty_list(self):\n        self.assertEqual(sorter([]), [])\n    def test_single_element_list(self):\n        self.assertEqual(sorter([5]), [5])\n    def test_ascending_order_list(self):\n        self.assertEqual(sorter([1, 2, 3, 4, 5]), [1, 2, 3, 4, 5])\n    def test_descending_order_list(self):\n        self.assertEqual(sorter([5, 4, 3, 2, 1]), [1, 2, 3, 4, 5])\n    def test_random_order_list(self):\n        self.assertEqual(sorter([3, 1, 4, 2, 5]), [1, 2, 3, 4, 5])\n    def test_duplicate_elements_list(self):\n        self.assertEqual(sorter([3, 1, 4, 2, 2, 5, 1]), [1, 1, 2, 2, 3, 4, 5])\n    def test_negative_numbers_list(self):\n        self.assertEqual(sorter([-5, -2, -8, -1, -3]), [-8, -5, -3, -2, -1])\n    def test_mixed_data_types_list(self):\n        self.assertEqual(sorter(['apple', 2, 'banana', 1, 'cherry']), [1, 2, 'apple', 'banana', 'cherry'])\n    def test_large_input_list(self):\n        self.assertEqual(sorter(list(range(1000, 0, -1))), list(range(1, 1001)))\n    def test_list_with_none_values(self):\n        self.assertEqual(sorter([None, 2, None, 1, None]), [None, None, None, 1, 2])\n    def test_list_with_nan_values(self):\n        self.assertEqual(sorter([float('nan'), 2, float('nan'), 1, float('nan')]), [1, 2, float('nan'), float('nan'), float('nan')])\n    def test_list_with_complex_numbers(self):\n        self.assertEqual(sorter([3 + 2j, 1 + 1j, 4 + 3j, 2 + 1j, 5 + 4j]), [1 + 1j, 2 + 1j, 3 + 2j, 4 + 3j, 5 + 4j])\n    def test_list_with_custom_class_objects(self):\n        class Person:\n            def __init__(self, name, age):\n                self.name = name\n                self.age = age\n            def __repr__(self):\n                return f\"Person('{self.name}', {self.age})\"\n        input_list = [Person('Alice', 25), Person('Bob', 30), Person('Charlie', 20)]\n        expected_output = [Person('Charlie', 20), Person('Alice', 25), Person('Bob', 30)]\n        self.assertEqual(sorter(input_list), expected_output)\n    def test_list_with_uncomparable_elements(self):\n        with self.assertRaises(TypeError):\n            sorter([5, 'apple', 3, [1, 2, 3], 2])\n    def test_list_with_custom_comparison_function(self):\n        input_list = [5, 4, 3, 2, 1]\n        expected_output = [5, 4, 3, 2, 1]\n        self.assertEqual(sorter(input_list, reverse=True), expected_output)\nif __name__ == '__main__':\n    unittest.main()"
+    return arr

code_to_optimize/bubble_sort_deps.py

@@ -9,142 +9,3 @@ def sorter_deps(arr):
                 dep2_swap(arr, j)
     return arr
 
-
-CACHED_TESTS = """import dill as pickle
-import os
-def _log__test__values(values, duration, test_name):
-    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
-    with open(os.path.join(
-    '/var/folders/ms/1tz2l1q55w5b7pp4wpdkbjq80000gn/T/codeflash_jk4pzz3w/', 
-    f'test_return_values_{iteration}.bin'), 'ab') as f:
-        return_bytes = pickle.dumps(values)
-        _test_name = f"{test_name}".encode("ascii")
-        f.write(len(_test_name).to_bytes(4, byteorder='big'))
-        f.write(_test_name)
-        f.write(duration.to_bytes(8, byteorder='big'))
-        f.write(len(return_bytes).to_bytes(4, byteorder='big'))
-        f.write(return_bytes)
-import time
-import gc
-from code_to_optimize.bubble_sort_deps import sorter_deps
-import timeout_decorator
-import unittest
-
-def dep1_comparer(arr, j: int) -> bool:
-    return arr[j] > arr[j + 1]
-
-def dep2_swap(arr, j):
-    temp = arr[j]
-    arr[j] = arr[j + 1]
-    arr[j + 1] = temp
-
-class TestSorterDeps(unittest.TestCase):
-
-    @timeout_decorator.timeout(15, use_signals=True)
-    def test_integers(self):
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([5, 3, 2, 4, 1])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(
-        return_value, duration, 
-        'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_integers:sorter_deps:0')
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([10, -3, 0, 2, 7])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(
-        return_value, duration, 
-        ('code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:'
-        'TestSorterDeps.test_integers:sorter_deps:1'))
-
-    @timeout_decorator.timeout(15, use_signals=True)
-    def test_floats(self):
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([3.2, 1.5, 2.7, 4.1, 1.0])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 
-        'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_floats:sorter_deps:0')
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([-1.1, 0.0, 3.14, 2.71, -0.5])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 
-        'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_floats:sorter_deps:1')
-
-    @timeout_decorator.timeout(15, use_signals=True)
-    def test_identical_elements(self):
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([1, 1, 1, 1, 1])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 
-        ('code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:'
-        'TestSorterDeps.test_identical_elements:sorter_deps:0'))
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([3.14, 3.14, 3.14])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 
-        ('code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:'
-        'TestSorterDeps.test_identical_elements:sorter_deps:1'))
-
-    @timeout_decorator.timeout(15, use_signals=True)
-    def test_single_element(self):
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([5])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_single_element:sorter_deps:0')
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([-3.2])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_single_element:sorter_deps:1')
-
-    @timeout_decorator.timeout(15, use_signals=True)
-    def test_empty_array(self):
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps([])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_empty_array:sorter_deps:0')
-
-    @timeout_decorator.timeout(15, use_signals=True)
-    def test_strings(self):
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps(['apple', 'banana', 'cherry', 'date'])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_strings:sorter_deps:0')
-        gc.disable()
-        counter = time.perf_counter_ns()
-        return_value = sorter_deps(['dog', 'cat', 'elephant', 'ant'])
-        duration = time.perf_counter_ns() - counter
-        gc.enable()
-        _log__test__values(return_value, duration, 'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_strings:sorter_deps:1')
-
-    @timeout_decorator.timeout(15, use_signals=True)
-    def test_mixed_types(self):
-        with self.assertRaises(TypeError):
-            gc.disable()
-            counter = time.perf_counter_ns()
-            return_value = sorter_deps([1, 'two', 3.0, 'four'])
-            duration = time.perf_counter_ns() - counter
-            gc.enable()
-            _log__test__values(return_value, duration, 'code_to_optimize.tests.unittest.test_sorter_deps__unit_test_0:TestSorterDeps.test_mixed_types:sorter_deps:0_0')
-if __name__ == '__main__':
-    unittest.main()
-
-"""

code_to_optimize/final_test_set/bubble_sort.py

@@ -5,7 +5,4 @@ def sorter(arr):
                 temp = arr[j]
                 arr[j] = arr[j + 1]
                 arr[j + 1] = temp
-    return arr
-
-
-CACHED_TESTS = "import unittest\ndef sorter(arr):\n    for i in range(len(arr)):\n        for j in range(len(arr) - 1):\n            if arr[j] > arr[j + 1]:\n                temp = arr[j]\n                arr[j] = arr[j + 1]\n                arr[j + 1] = temp\n    return arr\nclass SorterTestCase(unittest.TestCase):\n    def test_empty_list(self):\n        self.assertEqual(sorter([]), [])\n    def test_single_element_list(self):\n        self.assertEqual(sorter([5]), [5])\n    def test_ascending_order_list(self):\n        self.assertEqual(sorter([1, 2, 3, 4, 5]), [1, 2, 3, 4, 5])\n    def test_descending_order_list(self):\n        self.assertEqual(sorter([5, 4, 3, 2, 1]), [1, 2, 3, 4, 5])\n    def test_random_order_list(self):\n        self.assertEqual(sorter([3, 1, 4, 2, 5]), [1, 2, 3, 4, 5])\n    def test_duplicate_elements_list(self):\n        self.assertEqual(sorter([3, 1, 4, 2, 2, 5, 1]), [1, 1, 2, 2, 3, 4, 5])\n    def test_negative_numbers_list(self):\n        self.assertEqual(sorter([-5, -2, -8, -1, -3]), [-8, -5, -3, -2, -1])\n    def test_mixed_data_types_list(self):\n        self.assertEqual(sorter(['apple', 2, 'banana', 1, 'cherry']), [1, 2, 'apple', 'banana', 'cherry'])\n    def test_large_input_list(self):\n        self.assertEqual(sorter(list(range(1000, 0, -1))), list(range(1, 1001)))\n    def test_list_with_none_values(self):\n        self.assertEqual(sorter([None, 2, None, 1, None]), [None, None, None, 1, 2])\n    def test_list_with_nan_values(self):\n        self.assertEqual(sorter([float('nan'), 2, float('nan'), 1, float('nan')]), [1, 2, float('nan'), float('nan'), float('nan')])\n    def test_list_with_complex_numbers(self):\n        self.assertEqual(sorter([3 + 2j, 1 + 1j, 4 + 3j, 2 + 1j, 5 + 4j]), [1 + 1j, 2 + 1j, 3 + 2j, 4 + 3j, 5 + 4j])\n    def test_list_with_custom_class_objects(self):\n        class Person:\n            def __init__(self, name, age):\n                self.name = name\n                self.age = age\n            def __repr__(self):\n                return f\"Person('{self.name}', {self.age})\"\n        input_list = [Person('Alice', 25), Person('Bob', 30), Person('Charlie', 20)]\n        expected_output = [Person('Charlie', 20), Person('Alice', 25), Person('Bob', 30)]\n        self.assertEqual(sorter(input_list), expected_output)\n    def test_list_with_uncomparable_elements(self):\n        with self.assertRaises(TypeError):\n            sorter([5, 'apple', 3, [1, 2, 3], 2])\n    def test_list_with_custom_comparison_function(self):\n        input_list = [5, 4, 3, 2, 1]\n        expected_output = [5, 4, 3, 2, 1]\n        self.assertEqual(sorter(input_list, reverse=True), expected_output)\nif __name__ == '__main__':\n    unittest.main()"
+    return arr

code_to_optimize/use_cosine_similarity_from_other_file.py

@@ -9,110 +9,4 @@ def use_cosine_similarity(
     top_k: Optional[int] = 5,
     score_threshold: Optional[float] = None,
 ) -> Tuple[List[Tuple[int, int]], List[float]]:
-    return cosine_similarity_top_k(X, Y, top_k, score_threshold)
-
-
-CACHED_TESTS = """import unittest
-import numpy as np
-from sklearn.metrics.pairwise import cosine_similarity
-from typing import List, Optional, Tuple, Union
-Matrix = Union[List[List[float]], List[np.ndarray], np.ndarray]
-def cosine_similarity_top_k(X: Matrix, Y: Matrix, top_k: Optional[int]=5, score_threshold: Optional[float]=None) -> Tuple[List[Tuple[int, int]], List[float]]:
-    \"\"\"Row-wise cosine similarity with optional top-k and score threshold filtering.
-    Args:
-        X: Matrix.
-        Y: Matrix, same width as X.
-        top_k: Max number of results to return.
-        score_threshold: Minimum cosine similarity of results.
-    Returns:
-        Tuple of two lists. First contains two-tuples of indices (X_idx, Y_idx),
-            second contains corresponding cosine similarities.
-    \"\"\"
-    if len(X) == 0 or len(Y) == 0:
-        return ([], [])
-    score_array = cosine_similarity(X, Y)
-    sorted_idxs = score_array.flatten().argsort()[::-1]
-    top_k = top_k or len(sorted_idxs)
-    top_idxs = sorted_idxs[:top_k]
-    score_threshold = score_threshold or -1.0
-    top_idxs = top_idxs[score_array.flatten()[top_idxs] > score_threshold]
-    ret_idxs = [(x // score_array.shape[1], x % score_array.shape[1]) for x in top_idxs]
-    scores = score_array.flatten()[top_idxs].tolist()
-    return (ret_idxs, scores)
-def use_cosine_similarity(X: Matrix, Y: Matrix, top_k: Optional[int]=5, score_threshold: Optional[float]=None) -> Tuple[List[Tuple[int, int]], List[float]]:
-    return cosine_similarity_top_k(X, Y, top_k, score_threshold)
-class TestUseCosineSimilarity(unittest.TestCase):
-    def test_normal_scenario(self):
-        X = [[1, 2, 3], [4, 5, 6]]
-        Y = [[7, 8, 9], [10, 11, 12]]
-        result = use_cosine_similarity(X, Y, top_k=1, score_threshold=0.5)
-        self.assertEqual(result, ([(0, 1)], [0.9746318461970762]))
-    def test_edge_case_empty_matrices(self):
-        X = []
-        Y = []
-        result = use_cosine_similarity(X, Y)
-        self.assertEqual(result, ([], []))
-    def test_edge_case_different_widths(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5]]
-        with self.assertRaises(ValueError):
-            use_cosine_similarity(X, Y)
-    def test_edge_case_negative_top_k(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5, 6]]
-        with self.assertRaises(IndexError):
-            use_cosine_similarity(X, Y, top_k=-1)
-    def test_edge_case_zero_top_k(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5, 6]]
-        result = use_cosine_similarity(X, Y, top_k=0)
-        self.assertEqual(result, ([], []))
-    def test_edge_case_negative_score_threshold(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5, 6]]
-        result = use_cosine_similarity(X, Y, score_threshold=-1.0)
-        self.assertEqual(result, ([(0, 0)], [0.9746318461970762]))
-    def test_edge_case_large_score_threshold(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5, 6]]
-        result = use_cosine_similarity(X, Y, score_threshold=2.0)
-        self.assertEqual(result, ([], []))
-    def test_exceptional_case_non_matrix_X(self):
-        X = [1, 2, 3]
-        Y = [[4, 5, 6]]
-        with self.assertRaises(ValueError):
-            use_cosine_similarity(X, Y)
-    def test_exceptional_case_non_integer_top_k(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5, 6]]
-        with self.assertRaises(TypeError):
-            use_cosine_similarity(X, Y, top_k='5')
-    def test_exceptional_case_non_float_score_threshold(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5, 6]]
-        with self.assertRaises(TypeError):
-            use_cosine_similarity(X, Y, score_threshold='0.5')
-    def test_special_values_nan_in_matrices(self):
-        X = [[1, 2, np.nan]]
-        Y = [[4, 5, 6]]
-        with self.assertRaises(ValueError):
-            use_cosine_similarity(X, Y)
-    def test_special_values_none_top_k(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5, 6]]
-        result = use_cosine_similarity(X, Y, top_k=None)
-        self.assertEqual(result, ([(0, 0)], [0.9746318461970762]))
-    def test_special_values_none_score_threshold(self):
-        X = [[1, 2, 3]]
-        Y = [[4, 5, 6]]
-        result = use_cosine_similarity(X, Y, score_threshold=None)
-        self.assertEqual(result, ([(0, 0)], [0.9746318461970762]))
-    def test_large_inputs(self):
-        X = np.random.rand(1000, 1000)
-        Y = np.random.rand(1000, 1000)
-        result = use_cosine_similarity(X, Y, top_k=10, score_threshold=0.5)
-        self.assertEqual(len(result[0]), 10)
-        self.assertEqual(len(result[1]), 10)
-        self.assertTrue(all((score > 0.5 for score in result[1])))
-if __name__ == '__main__':
-    unittest.main()"""
+    return cosine_similarity_top_k(X, Y, top_k, score_threshold)

codeflash/code_utils/coverage_utils.py

@@ -12,7 +12,7 @@
 
 def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]:
     """Extract the single dependent function from the code context excluding the main function."""
-    ast_tree = ast.parse(code_context.code_to_optimize_with_helpers)
+    ast_tree = ast.parse(code_context.testgen_context_code)
 
     dependent_functions = {node.name for node in ast_tree.body if isinstance(node, ast.FunctionDef)}