refactoring

Author: aseembits93

codeflash/code_utils/code_utils.py

@@ -52,6 +52,39 @@ def custom_addopts() -> None:
                 f.write(original_content)
 
 
+@contextmanager
+def add_addopts_to_pyproject() -> None:
+    pyproject_file = find_pyproject_toml()
+    original_content = None
+    try:
+        # Read original file
+        if pyproject_file.exists():
+            with Path.open(pyproject_file, encoding="utf-8") as f:
+                original_content = f.read()
+                data = tomlkit.parse(original_content)
+            data["tool"]["pytest"] = {}
+            data["tool"]["pytest"]["ini_options"] = {}
+            data["tool"]["pytest"]["ini_options"]["addopts"] = [
+                "-n=auto",
+                "-n",
+                "1",
+                "-n 1",
+                "-n      1",
+                "-n      auto",
+            ]
+            with Path.open(
+                (Path(__file__).parent.parent.parent / "pyproject.toml").resolve(), "w", encoding="utf-8"
+            ) as f:
+                f.write(tomlkit.dumps(data))
+
+        yield
+
+    finally:
+        # Restore original file
+        with Path.open(Path("pyproject.toml"), "w", encoding="utf-8") as f:
+            f.write(original_content)
+
+
 def encoded_tokens_len(s: str) -> int:
     """Return the approximate length of the encoded tokens.
 

tests/scripts/end_to_end_test_topological_sort.py

@@ -2,24 +2,12 @@
 import pathlib
 import tomlkit
 
+from codeflash.code_utils.code_utils import add_addopts_to_pyproject
 from end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command, run_with_retries
 
 
 def run_test(expected_improvement_pct: int) -> bool:
-    try:
-        # Modify Pyproject file
-        with pathlib.Path.open(
-            (pathlib.Path(__file__).parent.parent.parent / "pyproject.toml").resolve(), encoding="utf-8"
-        ) as f:
-            original_content = f.read()
-            data = tomlkit.parse(original_content)
-        data["tool"]["pytest"] = {}
-        data["tool"]["pytest"]["ini_options"] = {}
-        data["tool"]["pytest"]["ini_options"]["addopts"] = ["-n=auto", "-n", "1", "-n 1", "-n      1", "-n      auto"]
-        with pathlib.Path.open(
-            (pathlib.Path(__file__).parent.parent.parent / "pyproject.toml").resolve(), "w", encoding="utf-8"
-        ) as f:
-            f.write(tomlkit.dumps(data))
+    with add_addopts_to_pyproject():
         config = TestConfig(
             file_path="topological_sort.py",
             function_name="Graph.topologicalSort",
@@ -35,9 +23,6 @@ def run_test(expected_improvement_pct: int) -> bool:
         )
         cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve()
         return_var = run_codeflash_command(cwd, config, expected_improvement_pct)
-    finally:
-        with pathlib.Path.open(pathlib.Path("pyproject.toml"), "w", encoding="utf-8") as f:
-            f.write(original_content)
     return return_var