Merge remote-tracking branch 'origin/main' into rewrite-candidate-loop-clean

Author: aseembits93

README.md

@@ -3,9 +3,7 @@
    <a href="https://github.com/codeflash-ai/codeflash">
     <img src="https://img.shields.io/github/commit-activity/m/codeflash-ai/codeflash" alt="GitHub commit activity">
   </a>
-  <a href="https://pypi.org/project/codeflash/">
-    <img src="https://img.shields.io/pypi/dm/codeflash" alt="PyPI Downloads">
-  </a>
+  <a href="https://pypi.org/project/codeflash/"><img src="https://static.pepy.tech/badge/codeflash" alt="PyPI Downloads"></a>
   <a href="https://pypi.org/project/codeflash/">
     <img src="https://img.shields.io/pypi/v/codeflash?label=PyPI%20version" alt="PyPI Downloads">
   </a>
@@ -83,4 +81,4 @@ Join our community for support and discussions. If you have any questions, feel
 
 ## License
 
-Codeflash is licensed under the BSL-1.1 License. See the LICENSE file for details.
+Codeflash is licensed under the BSL-1.1 License. See the [LICENSE](https://github.com/codeflash-ai/codeflash/blob/main/codeflash/LICENSE) file for details.

codeflash/LICENSE

@@ -3,7 +3,7 @@ Business Source License 1.1
 Parameters
 
 Licensor:             CodeFlash Inc.
-Licensed Work:        Codeflash Client version 0.15.x
+Licensed Work:        Codeflash Client version 0.16.x
                       The Licensed Work is (c) 2024 CodeFlash Inc.
 
 Additional Use Grant: None. Production use of the Licensed Work is only permitted
@@ -13,7 +13,7 @@ Additional Use Grant: None. Production use of the Licensed Work is only permitte
                       Platform. Please visit codeflash.ai for further
                       information.
 
-Change Date:          2029-07-03
+Change Date:          2029-08-14
 
 Change License:       MIT
 

codeflash/api/aiservice.py

@@ -10,8 +10,9 @@
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
-from codeflash.code_utils.env_utils import get_codeflash_api_key, is_LSP_enabled
+from codeflash.code_utils.env_utils import get_codeflash_api_key
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
+from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import AIServiceRefinerRequest, CodeStringsMarkdown, OptimizedCandidate
 from codeflash.telemetry.posthog_cf import ph

codeflash/api/cfapi.py

@@ -14,8 +14,9 @@
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
-from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name, git_root_dir
+from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name
 from codeflash.github.PrComment import FileDiffContent, PrComment
+from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.version import __version__
 
 if TYPE_CHECKING:
@@ -101,7 +102,7 @@ def get_user_id() -> Optional[str]:
             if min_version and version.parse(min_version) > version.parse(__version__):
                 msg = "Your Codeflash CLI version is outdated. Please update to the latest version using `pip install --upgrade codeflash`."
                 console.print(f"[bold red]{msg}[/bold red]")
-                if console.quiet:  # lsp
+                if is_LSP_enabled():
                     logger.debug(msg)
                     return f"Error: {msg}"
                 sys.exit(1)
@@ -203,8 +204,9 @@ def create_staging(
     generated_original_test_source: str,
     function_trace_id: str,
     coverage_message: str,
-    replay_tests: str = "",
-    concolic_tests: str = "",
+    replay_tests: str,
+    concolic_tests: str,
+    root_dir: Path,
 ) -> Response:
     """Create a staging pull request, targeting the specified branch. (usually 'staging').
 
@@ -217,12 +219,10 @@ def create_staging(
     :param coverage_message: Coverage report or summary.
     :return: The response object from the backend.
     """
-    relative_path = explanation.file_path.relative_to(git_root_dir()).as_posix()
+    relative_path = explanation.file_path.relative_to(root_dir).as_posix()
 
     build_file_changes = {
-        Path(p).relative_to(git_root_dir()).as_posix(): FileDiffContent(
-            oldContent=original_code[p], newContent=new_code[p]
-        )
+        Path(p).relative_to(root_dir).as_posix(): FileDiffContent(oldContent=original_code[p], newContent=new_code[p])
         for p in original_code
     }
 

codeflash/cli_cmds/cli.py

@@ -10,6 +10,7 @@
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.config_parser import parse_config_file
+from codeflash.lsp.helpers import is_LSP_enabled
 from codeflash.version import __version__ as version
 
 
@@ -94,6 +95,7 @@ def parse_args() -> Namespace:
         help="Path to the directory of the project, where all the pytest-benchmark tests are located.",
     )
     parser.add_argument("--no-draft", default=False, action="store_true", help="Skip optimization for draft PRs")
+    parser.add_argument("--worktree", default=False, action="store_true", help="Use worktree for optimization")
 
     args, unknown_args = parser.parse_known_args()
     sys.argv[:] = [sys.argv[0], *unknown_args]
@@ -210,6 +212,9 @@ def process_pyproject_config(args: Namespace) -> Namespace:
     if args.benchmarks_root:
         args.benchmarks_root = Path(args.benchmarks_root).resolve()
     args.test_project_root = project_root_from_module_root(args.tests_root, pyproject_file_path)
+    if is_LSP_enabled():
+        args.all = None
+        return args
     return handle_optimize_all_arg_parsing(args)
 
 

codeflash/cli_cmds/cmd_init.py

@@ -155,6 +155,22 @@ def ask_run_end_to_end_test(args: Namespace) -> None:
         run_end_to_end_test(args, bubble_sort_path, bubble_sort_test_path)
 
 
+def is_valid_pyproject_toml(pyproject_toml_path: Path) -> dict[str, Any] | None:
+    if not pyproject_toml_path.exists():
+        return None
+    try:
+        config, _ = parse_config_file(pyproject_toml_path)
+    except Exception:
+        return None
+
+    if "module_root" not in config or config["module_root"] is None or not Path(config["module_root"]).is_dir():
+        return None
+    if "tests_root" not in config or config["tests_root"] is None or not Path(config["tests_root"]).is_dir():
+        return None
+
+    return config
+
+
 def should_modify_pyproject_toml() -> tuple[bool, dict[str, Any] | None]:
     """Check if the current directory contains a valid pyproject.toml file with codeflash config.
 
@@ -163,16 +179,9 @@ def should_modify_pyproject_toml() -> tuple[bool, dict[str, Any] | None]:
     from rich.prompt import Confirm
 
     pyproject_toml_path = Path.cwd() / "pyproject.toml"
-    if not pyproject_toml_path.exists():
-        return True, None
-    try:
-        config, config_file_path = parse_config_file(pyproject_toml_path)
-    except Exception:
-        return True, None
 
-    if "module_root" not in config or config["module_root"] is None or not Path(config["module_root"]).is_dir():
-        return True, None
-    if "tests_root" not in config or config["tests_root"] is None or not Path(config["tests_root"]).is_dir():
+    config = is_valid_pyproject_toml(pyproject_toml_path)
+    if config is None:
         return True, None
 
     return Confirm.ask(
@@ -968,6 +977,11 @@ def install_github_app(git_remote: str) -> None:
     except git.InvalidGitRepositoryError:
         click.echo("Skipping GitHub app installation because you're not in a git repository.")
         return
+
+    if git_remote not in get_git_remotes(git_repo):
+        click.echo(f"Skipping GitHub app installation, remote ({git_remote}) does not exist in this repository.")
+        return
+
     owner, repo = get_repo_owner_and_name(git_repo, git_remote)
 
     if is_github_app_installed_on_repo(owner, repo, suppress_errors=True):

codeflash/cli_cmds/console.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from contextlib import contextmanager
 from itertools import cycle
 from typing import TYPE_CHECKING
@@ -28,6 +29,10 @@
 DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG
 
 console = Console()
+
+if os.getenv("CODEFLASH_LSP"):
+    console.quiet = True
+
 logging.basicConfig(
     level=logging.INFO,
     handlers=[RichHandler(rich_tracebacks=True, markup=False, console=console, show_path=False, show_time=False)],

codeflash/code_utils/code_extractor.py

@@ -195,8 +195,19 @@ def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None:
                 self.last_import_line = self.current_line
 
 
-class ConditionalImportCollector(cst.CSTVisitor):
-    """Collect imports inside top-level conditionals (e.g., if TYPE_CHECKING, try/except)."""
+class DottedImportCollector(cst.CSTVisitor):
+    """Collects all top-level imports from a Python module in normalized dotted format, including top-level conditional imports like `if TYPE_CHECKING:`.
+
+    Examples
+    --------
+        import os                                                                  ==> "os"
+        import dbt.adapters.factory                                                ==> "dbt.adapters.factory"
+        from pathlib import Path                                                   ==> "pathlib.Path"
+        from recce.adapter.base import BaseAdapter                                 ==> "recce.adapter.base.BaseAdapter"
+        from typing import Any, List, Optional                                     ==> "typing.Any", "typing.List", "typing.Optional"
+        from recce.util.lineage import ( build_column_key, filter_dependency_maps) ==> "recce.util.lineage.build_column_key", "recce.util.lineage.filter_dependency_maps"
+
+    """
 
     def __init__(self) -> None:
         self.imports: set[str] = set()
@@ -217,7 +228,10 @@ def _collect_imports_from_block(self, block: cst.IndentedBlock) -> None:
                         for alias in child.names:
                             module = self.get_full_dotted_name(alias.name)
                             asname = alias.asname.name.value if alias.asname else alias.name.value
-                            self.imports.add(module if module == asname else f"{module}.{asname}")
+                            if isinstance(asname, cst.Attribute):
+                                self.imports.add(module)
+                            else:
+                                self.imports.add(module if module == asname else f"{module}.{asname}")
 
                     elif isinstance(child, cst.ImportFrom):
                         if child.module is None:
@@ -231,6 +245,7 @@ def _collect_imports_from_block(self, block: cst.IndentedBlock) -> None:
 
     def visit_Module(self, node: cst.Module) -> None:
         self.depth = 0
+        self._collect_imports_from_block(node)
 
     def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
         self.depth += 1
@@ -388,45 +403,44 @@ def add_needed_imports_from_module(
         logger.error(f"Error parsing source module code: {e}")
         return dst_module_code
 
-    cond_import_collector = ConditionalImportCollector()
+    dotted_import_collector = DottedImportCollector()
     try:
         parsed_dst_module = cst.parse_module(dst_module_code)
-        parsed_dst_module.visit(cond_import_collector)
+        parsed_dst_module.visit(dotted_import_collector)
     except cst.ParserSyntaxError as e:
         logger.exception(f"Syntax error in destination module code: {e}")
         return dst_module_code  # Return the original code if there's a syntax error
 
     try:
         for mod in gatherer.module_imports:
-            if mod in cond_import_collector.imports:
-                continue
-            AddImportsVisitor.add_needed_import(dst_context, mod)
+            if mod not in dotted_import_collector.imports:
+                AddImportsVisitor.add_needed_import(dst_context, mod)
             RemoveImportsVisitor.remove_unused_import(dst_context, mod)
         for mod, obj_seq in gatherer.object_mapping.items():
             for obj in obj_seq:
                 if (
                     f"{mod}.{obj}" in helper_functions_fqn or dst_context.full_module_name == mod  # avoid circular deps
                 ):
                     continue  # Skip adding imports for helper functions already in the context
-                if f"{mod}.{obj}" in cond_import_collector.imports:
-                    continue
-                AddImportsVisitor.add_needed_import(dst_context, mod, obj)
+                if f"{mod}.{obj}" not in dotted_import_collector.imports:
+                    AddImportsVisitor.add_needed_import(dst_context, mod, obj)
                 RemoveImportsVisitor.remove_unused_import(dst_context, mod, obj)
     except Exception as e:
         logger.exception(f"Error adding imports to destination module code: {e}")
         return dst_module_code
+
     for mod, asname in gatherer.module_aliases.items():
-        if f"{mod}.{asname}" in cond_import_collector.imports:
-            continue
-        AddImportsVisitor.add_needed_import(dst_context, mod, asname=asname)
+        if f"{mod}.{asname}" not in dotted_import_collector.imports:
+            AddImportsVisitor.add_needed_import(dst_context, mod, asname=asname)
         RemoveImportsVisitor.remove_unused_import(dst_context, mod, asname=asname)
+
     for mod, alias_pairs in gatherer.alias_mapping.items():
         for alias_pair in alias_pairs:
             if f"{mod}.{alias_pair[0]}" in helper_functions_fqn:
                 continue
-            if f"{mod}.{alias_pair[1]}" in cond_import_collector.imports:
-                continue
-            AddImportsVisitor.add_needed_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
+
+            if f"{mod}.{alias_pair[1]}" not in dotted_import_collector.imports:
+                AddImportsVisitor.add_needed_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
             RemoveImportsVisitor.remove_unused_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
 
     try:

codeflash/code_utils/coverage_utils.py

@@ -39,16 +39,20 @@ def build_fully_qualified_name(function_name: str, code_context: CodeOptimizatio
     return full_name
 
 
-def generate_candidates(source_code_path: Path) -> list[str]:
+def generate_candidates(source_code_path: Path) -> set[str]:
     """Generate all the possible candidates for coverage data based on the source code path."""
-    candidates = [source_code_path.name]
+    candidates = set()
+    candidates.add(source_code_path.name)
     current_path = source_code_path.parent
 
+    last_added = source_code_path.name
     while current_path != current_path.parent:
-        candidate_path = str(Path(current_path.name) / candidates[-1])
-        candidates.append(candidate_path)
+        candidate_path = str(Path(current_path.name) / last_added)
+        candidates.add(candidate_path)
+        last_added = candidate_path
         current_path = current_path.parent
 
+    candidates.add(str(source_code_path))
     return candidates
 
 

codeflash/code_utils/env_utils.py

@@ -7,10 +7,11 @@
 from pathlib import Path
 from typing import Any, Optional
 
-from codeflash.cli_cmds.console import console, logger
+from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import exit_with_message
 from codeflash.code_utils.formatter import format_code
 from codeflash.code_utils.shell_utils import read_api_key_from_shell_config
+from codeflash.lsp.helpers import is_LSP_enabled
 
 
 def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool = True) -> bool:  # noqa
@@ -34,11 +35,12 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool =
 
 @lru_cache(maxsize=1)
 def get_codeflash_api_key() -> str:
-    if console.quiet:  # lsp
-        # prefer shell config over env var in lsp mode
-        api_key = read_api_key_from_shell_config()
-    else:
-        api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+    # prefer shell config over env var in lsp mode
+    api_key = (
+        read_api_key_from_shell_config()
+        if is_LSP_enabled()
+        else os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()
+    )
 
     api_secret_docs_message = "For more information, refer to the documentation at [https://docs.codeflash.ai/getting-started/codeflash-github-actions#add-your-api-key-to-your-repository-secrets]."  # noqa
     if not api_key:
@@ -125,11 +127,6 @@ def is_ci() -> bool:
     return bool(os.environ.get("CI") or os.environ.get("GITHUB_ACTIONS"))
 
 
-@lru_cache(maxsize=1)
-def is_LSP_enabled() -> bool:
-    return console.quiet
-
-
 def is_pr_draft() -> bool:
     """Check if the PR is draft. in the github action context."""
     event = get_cached_gh_event_data()