Add precautions to FD leaks2

Author: Saga4

codeflash/benchmarking/replay_test.py

@@ -24,22 +24,25 @@ def get_next_arg_and_return(
     num_to_get: int = 256,
 ) -> Generator[Any]:
     db = sqlite3.connect(trace_file)
-    cur = db.cursor()
-    limit = num_to_get
-
-    if class_name is not None:
-        cursor = cur.execute(
-            "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = ? LIMIT ?",
-            (benchmark_function_name, function_name, file_path, class_name, limit),
-        )
-    else:
-        cursor = cur.execute(
-            "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = '' LIMIT ?",
-            (benchmark_function_name, function_name, file_path, limit),
-        )
+    try:
+        cur = db.cursor()
+        limit = num_to_get
+
+        if class_name is not None:
+            cursor = cur.execute(
+                "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = ? LIMIT ?",
+                (benchmark_function_name, function_name, file_path, class_name, limit),
+            )
+        else:
+            cursor = cur.execute(
+                "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = '' LIMIT ?",
+                (benchmark_function_name, function_name, file_path, limit),
+            )
 
-    while (val := cursor.fetchone()) is not None:
-        yield val[9], val[10]  # pickled_args, pickled_kwargs
+        while (val := cursor.fetchone()) is not None:
+            yield val[9], val[10]  # pickled_args, pickled_kwargs
+    finally:
+        db.close()
 
 
 def get_function_alias(module: str, function_name: str) -> str:
@@ -235,61 +238,69 @@ def generate_replay_test(
     try:
         # Connect to the database
         conn = sqlite3.connect(trace_file_path.as_posix())
-        cursor = conn.cursor()
-
-        # Get distinct benchmark file paths
-        cursor.execute("SELECT DISTINCT benchmark_module_path FROM benchmark_function_timings")
-        benchmark_files = cursor.fetchall()
-
-        # Generate a test for each benchmark file
-        for benchmark_file in benchmark_files:
-            benchmark_module_path = benchmark_file[0]
-            # Get all benchmarks and functions associated with this file path
-            cursor.execute(
-                "SELECT DISTINCT benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings "
-                "WHERE benchmark_module_path = ?",
-                (benchmark_module_path,),
-            )
-
-            functions_data = []
-            for row in cursor.fetchall():
-                benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number = row
-                # Add this function to our list
-                functions_data.append(
-                    {
-                        "function_name": function_name,
-                        "class_name": class_name,
-                        "file_path": file_path,
-                        "module_name": module_name,
-                        "benchmark_function_name": benchmark_function_name,
-                        "benchmark_module_path": benchmark_module_path,
-                        "benchmark_line_number": benchmark_line_number,
-                        "function_properties": inspect_top_level_functions_or_methods(
-                            file_name=Path(file_path), function_or_method_name=function_name, class_name=class_name
-                        ),
-                    }
+        try:
+            cursor = conn.cursor()
+
+            # Get distinct benchmark file paths
+            cursor.execute("SELECT DISTINCT benchmark_module_path FROM benchmark_function_timings")
+            benchmark_files = cursor.fetchall()
+
+            # Generate a test for each benchmark file
+            for benchmark_file in benchmark_files:
+                benchmark_module_path = benchmark_file[0]
+                # Get all benchmarks and functions associated with this file path
+                cursor.execute(
+                    "SELECT DISTINCT benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings "
+                    "WHERE benchmark_module_path = ?",
+                    (benchmark_module_path,),
                 )
 
-            if not functions_data:
-                logger.info(f"No benchmark test functions found in {benchmark_module_path}")
-                continue
-            # Generate the test code for this benchmark
-            test_code = create_trace_replay_test_code(
-                trace_file=trace_file_path.as_posix(),
-                functions_data=functions_data,
-                test_framework=test_framework,
-                max_run_count=max_run_count,
-            )
-            test_code = isort.code(test_code)
-            output_file = get_test_file_path(
-                test_dir=Path(output_dir), function_name=benchmark_module_path, test_type="replay"
-            )
-            # Write test code to file, parents = true
-            output_dir.mkdir(parents=True, exist_ok=True)
-            output_file.write_text(test_code, "utf-8")
-            count += 1
-
-        conn.close()
+                functions_data = []
+                for row in cursor.fetchall():
+                    (
+                        benchmark_function_name,
+                        function_name,
+                        class_name,
+                        module_name,
+                        file_path,
+                        benchmark_line_number,
+                    ) = row
+                    # Add this function to our list
+                    functions_data.append(
+                        {
+                            "function_name": function_name,
+                            "class_name": class_name,
+                            "file_path": file_path,
+                            "module_name": module_name,
+                            "benchmark_function_name": benchmark_function_name,
+                            "benchmark_module_path": benchmark_module_path,
+                            "benchmark_line_number": benchmark_line_number,
+                            "function_properties": inspect_top_level_functions_or_methods(
+                                file_name=Path(file_path), function_or_method_name=function_name, class_name=class_name
+                            ),
+                        }
+                    )
+
+                if not functions_data:
+                    logger.info(f"No benchmark test functions found in {benchmark_module_path}")
+                    continue
+                # Generate the test code for this benchmark
+                test_code = create_trace_replay_test_code(
+                    trace_file=trace_file_path.as_posix(),
+                    functions_data=functions_data,
+                    test_framework=test_framework,
+                    max_run_count=max_run_count,
+                )
+                test_code = isort.code(test_code)
+                output_file = get_test_file_path(
+                    test_dir=Path(output_dir), function_name=benchmark_module_path, test_type="replay"
+                )
+                # Write test code to file, parents = true
+                output_dir.mkdir(parents=True, exist_ok=True)
+                output_file.write_text(test_code, "utf-8")
+                count += 1
+        finally:
+            conn.close()
     except Exception as e:
         logger.info(f"Error generating replay tests: {e}")
 

codeflash/lsp/beta.py

@@ -109,8 +109,7 @@ def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOpt
     fto = server.optimizer.current_function_being_optimized
     optimizable_funcs = {fto.file_path: [fto]}
 
-    devnull_writer = open(os.devnull, "w")  # noqa
-    with contextlib.redirect_stdout(devnull_writer):
+    with open(os.devnull, "w") as devnull_writer, contextlib.redirect_stdout(devnull_writer):  # noqa: PTH123
         function_to_tests, num_discovered_tests = server.optimizer.discover_tests(optimizable_funcs)
 
     server.optimizer.discovered_tests = function_to_tests