Merge branch 'refs/heads/main' into testgen-context-improvement

# Conflicts:
#	codeflash/context/code_context_extractor.py
#	tests/test_get_read_writable_code.py

Author: alvin-r

.github/workflows/pr_agent.yml

@@ -0,0 +1,20 @@
+on:
+  pull_request:
+    types: [opened, reopened, ready_for_review]
+  issue_comment:
+jobs:
+  pr_agent_job:
+    if: ${{ github.event.sender.type != 'Bot' }}
+    runs-on: ubuntu-latest
+    permissions:
+      issues: write
+      pull-requests: write
+      contents: write
+    name: Run pr agent on every pull request, respond to user comments
+    steps:
+      - name: PR Agent action step
+        id: pragent
+        uses: qodo-ai/pr-agent@main
+        env:
+          OPENAI_KEY: ${{ secrets.OPENAI_KEY }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

code_to_optimize/bubble_sort.py

@@ -1,8 +1,10 @@
 def sorter(arr):
+    print("codeflash stdout: Sorting list")
     for i in range(len(arr)):
         for j in range(len(arr) - 1):
             if arr[j] > arr[j + 1]:
                 temp = arr[j]
                 arr[j] = arr[j + 1]
                 arr[j + 1] = temp
+    print(f"result: {arr}")
     return arr

code_to_optimize/bubble_sort_method.py

@@ -1,12 +1,17 @@
+import sys
+
+
 class BubbleSorter:
     def __init__(self, x=0):
         self.x = x
 
     def sorter(self, arr):
+        print("codeflash stdout : BubbleSorter.sorter() called")
         for i in range(len(arr)):
             for j in range(len(arr) - 1):
                 if arr[j] > arr[j + 1]:
                     temp = arr[j]
                     arr[j] = arr[j + 1]
                     arr[j + 1] = temp
+        print("stderr test", file=sys.stderr)
         return arr

codeflash.code-workspace

@@ -0,0 +1,80 @@
+{
+    "folders": [
+        {
+            "path": ".",
+            "name": "codeflash",
+            "extensions": [
+                "charliermarsh.ruff",
+                "ms-python.python",
+            ]
+        }
+    ],
+    "settings": {
+        "python.defaultInterpreterPath": "~/miniforge3/envs/codeflash312/bin/python",
+        "python.terminal.activateEnvironment": true,
+        "python.testing.pytestEnabled": true,
+        "python.testing.pytestArgs": ["tests/", "-vv"],
+    },
+    "launch": {
+        "version": "0.2.0",
+        "configurations": [
+            {
+                "name": "bubble_sort",
+                "type": "debugpy",
+                "request": "launch",
+                "program": "${workspaceFolder:codeflash}/codeflash/main.py",
+                "args": [
+                    "--file",
+                    "code_to_optimize/bubble_sort.py",
+                    "--module-root",
+                    "${workspaceFolder:codeflash}",
+                    "--function",
+                    "sorter",
+                    "--test-framework",
+                    "pytest",
+                    "--tests-root",
+                    "code_to_optimize/tests/pytest"
+                ],
+                "cwd": "${workspaceFolder:codeflash}",
+                "console": "integratedTerminal",
+                "env": {
+                    "PYTHONUNBUFFERED": "1"
+                },
+            },
+            {
+                "name": "bubble_sort -all",
+                "type": "debugpy",
+                "request": "launch",
+                "program": "${workspaceFolder:codeflash}/codeflash/main.py",
+                "args": [
+                    "--all",
+                    "--test-framework",
+                    "pytest",
+                    "--tests-root",
+                    "code_to_optimize/tests/pytest",
+                    "--module-root",
+                    "code_to_optimize"
+                ],
+                "cwd": "${workspaceFolder:codeflash}",
+                "console": "integratedTerminal",
+                "env": {
+                    "PYTHONUNBUFFERED": "1"
+                },
+            },
+            {
+                "name": "bubble_sort --file bubble_sort.py (MBR)",
+                "type": "debugpy",
+                "request": "launch",
+                "program": "${workspaceFolder:codeflash}/codeflash/main.py",
+                "args": [
+                    "--all",
+                ],
+                "cwd": "/Users/krrt7/Desktop/work/my-best-repo",
+                "console": "integratedTerminal",
+                "env": {
+                    "PYTHONUNBUFFERED": "1"
+                },
+            }
+        ]
+    }
+}

codeflash/api/aiservice.py

@@ -29,6 +29,7 @@ def __init__(self) -> None:
     def get_aiservice_base_url(self) -> str:
         if os.environ.get("CODEFLASH_AIS_SERVER", default="prod").lower() == "local":
             logger.info("Using local AI Service at http://localhost:8000")
+            console.rule()
             return "http://localhost:8000"
         return "https://app.codeflash.ai"
 
@@ -196,10 +197,9 @@ def generate_regression_tests(
         - Dict[str, str] | None: The generated regression tests and instrumented tests, or None if an error occurred.
 
         """
-        assert test_framework in [
-            "pytest",
-            "unittest",
-        ], f"Invalid test framework, got {test_framework} but expected 'pytest' or 'unittest'"
+        assert test_framework in ["pytest", "unittest"], (
+            f"Invalid test framework, got {test_framework} but expected 'pytest' or 'unittest'"
+        )
         payload = {
             "source_code_being_tested": source_code_being_tested,
             "function_to_optimize": function_to_optimize,

codeflash/cli_cmds/cmd_init.py

@@ -462,7 +462,7 @@ def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyMa
         return DependencyManager.POETRY
 
     # Check for uv
-    if any(key.startswith("uv") for key in tool_section.keys()):
+    if any(key.startswith("uv") for key in tool_section):
         return DependencyManager.UV
 
     # Look for pip-specific markers
@@ -555,9 +555,7 @@ def customize_codeflash_yaml_content(
 
     # Add codeflash command
     codeflash_cmd = get_codeflash_github_action_command(dep_manager)
-    optimize_yml_content = optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd)
-
-    return optimize_yml_content
+    return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd)
 
 
 # Create or update the pyproject.toml file with the Codeflash dependency & configuration
@@ -596,8 +594,8 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
         formatter_cmds.append("disabled")
     if formatter in ["black", "ruff"]:
         try:
-            result = subprocess.run([formatter], capture_output=True, check=False)
-        except FileNotFoundError as e:
+            subprocess.run([formatter], capture_output=True, check=False)
+        except FileNotFoundError:
             click.echo(f"⚠️ Formatter not found: {formatter}, please ensure it is installed")
     codeflash_section["formatter-cmds"] = formatter_cmds
     # Add the 'codeflash' section, ensuring 'tool' section exists

codeflash/cli_cmds/console.py

@@ -3,7 +3,7 @@
 import logging
 from contextlib import contextmanager
 from itertools import cycle
-from typing import TYPE_CHECKING, Generator
+from typing import TYPE_CHECKING
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -13,6 +13,8 @@
 from codeflash.cli_cmds.logging_config import BARE_LOGGING_FORMAT
 
 if TYPE_CHECKING:
+    from collections.abc import Generator
+
     from rich.progress import TaskID
 
 DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG

codeflash/code_utils/code_extractor.py

@@ -115,8 +115,7 @@ def get_code(functions_to_optimize: list[FunctionToOptimize]) -> tuple[str | Non
         or (functions_to_optimize[0].parents and functions_to_optimize[0].parents[0].type != "ClassDef")
         or (
             len(functions_to_optimize[0].parents) > 1
-            or (len(functions_to_optimize) > 1)
-            and len({fn.parents[0] for fn in functions_to_optimize}) != 1
+            or ((len(functions_to_optimize) > 1) and len({fn.parents[0] for fn in functions_to_optimize}) != 1)
         )
     ):
         return None, set()

codeflash/code_utils/code_replacer.py

@@ -3,7 +3,7 @@
 import ast
 from collections import defaultdict
 from functools import lru_cache
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Optional, TypeVar
 
 import libcst as cst
 
@@ -90,10 +90,10 @@ def leave_ClassDef(self, node: cst.ClassDef) -> None:
 class OptimFunctionReplacer(cst.CSTTransformer):
     def __init__(
         self,
-        modified_functions: dict[tuple[str | None, str], cst.FunctionDef] = None,
-        new_functions: list[cst.FunctionDef] = None,
-        new_class_functions: dict[str, list[cst.FunctionDef]] = None,
-        modified_init_functions: dict[str, cst.FunctionDef] = None,
+        modified_functions: Optional[dict[tuple[str | None, str], cst.FunctionDef]] = None,
+        new_functions: Optional[list[cst.FunctionDef]] = None,
+        new_class_functions: Optional[dict[str, list[cst.FunctionDef]]] = None,
+        modified_init_functions: Optional[dict[str, cst.FunctionDef]] = None,
     ) -> None:
         super().__init__()
         self.modified_functions = modified_functions if modified_functions is not None else {}

codeflash/code_utils/code_utils.py

@@ -12,7 +12,8 @@
 
 def get_qualified_name(module_name: str, full_qualified_name: str) -> str:
     if not full_qualified_name:
-        raise ValueError("full_qualified_name cannot be empty")
+        msg = "full_qualified_name cannot be empty"
+        raise ValueError(msg)
     if not full_qualified_name.startswith(module_name):
         msg = f"{full_qualified_name} does not start with {module_name}"
         raise ValueError(msg)
@@ -46,9 +47,9 @@ def file_name_from_test_module_name(test_module_name: str, base_dir: Path) -> Pa
 def get_imports_from_file(
     file_path: Path | None = None, file_string: str | None = None, file_ast: ast.AST | None = None
 ) -> list[ast.Import | ast.ImportFrom]:
-    assert (
-        sum([file_path is not None, file_string is not None, file_ast is not None]) == 1
-    ), "Must provide exactly one of file_path, file_string, or file_ast"
+    assert sum([file_path is not None, file_string is not None, file_ast is not None]) == 1, (
+        "Must provide exactly one of file_path, file_string, or file_ast"
+    )
     if file_path:
         with file_path.open(encoding="utf8") as file:
             file_string = file.read()
@@ -107,6 +108,14 @@ def validate_python_code(code: str) -> str:
     return code
 
 
+def has_any_async_functions(code: str) -> bool:
+    try:
+        module = ast.parse(code)
+    except SyntaxError:
+        return False
+    return any(isinstance(node, ast.AsyncFunctionDef) for node in ast.walk(module))
+
+
 def cleanup_paths(paths: list[Path]) -> None:
     for path in paths:
         path.unlink(missing_ok=True)