From f43ee06859417d09f09e7c5a9fff8896ae99265b Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Sat, 7 Mar 2026 04:42:50 -0500 Subject: [PATCH] refactor: restructure codebase for locality and faster CLI startup Move files closer to their consumers: - function_context.py merged into code_context_extractor.py - FunctionOptimizer base class to languages/function_optimizer.py - test_runner, instrument_codeflash_capture, parse_line_profile to languages/python/ - oauth_handler.py to cli_cmds/ Split cmd_init.py (1993 lines) into focused modules: - init_config.py: config types, validation, writing, shared UI - init_auth.py: API key management + GitHub app installation - github_workflow.py: GitHub Actions workflow generation - cmd_init.py: init orchestrator + Python setup (639 lines) Defer heavy imports (cmd_init, posthog, sentry) from module-level to usage sites, reducing CLI startup from ~600ms to ~250ms. Replace set_defaults(func=) with direct args.command dispatch in main(). --- .claude/rules/architecture.md | 19 +- codeflash/cli_cmds/cli.py | 21 +- codeflash/cli_cmds/cmd_init.py | 1403 +---------------- codeflash/cli_cmds/console_constants.py | 15 + codeflash/cli_cmds/github_workflow.py | 898 +++++++++++ codeflash/cli_cmds/init_auth.py | 209 +++ codeflash/cli_cmds/init_config.py | 286 ++++ codeflash/cli_cmds/init_javascript.py | 7 +- .../{code_utils => cli_cmds}/oauth_handler.py | 0 codeflash/languages/base.py | 2 +- .../function_optimizer.py | 0 .../javascript/function_optimizer.py | 2 +- .../python/context/code_context_extractor.py | 39 +- .../languages/python/function_optimizer.py | 6 +- .../python}/instrument_codeflash_capture.py | 0 .../python}/parse_line_profile_test_output.py | 0 codeflash/languages/python/reference_graph.py | 2 +- codeflash/languages/python/support.py | 6 +- .../python}/test_runner.py | 0 codeflash/lsp/beta.py | 4 +- codeflash/main.py | 29 +- codeflash/optimization/function_context.py | 46 - codeflash/optimization/optimizer.py | 2 +- codeflash/tracer.py | 2 +- ...est_benchmark_code_extract_code_context.py | 2 +- tests/test_async_run_and_parse_tests.py | 2 +- tests/test_cmd_init.py | 2 +- tests/test_codeflash_capture.py | 4 +- tests/test_formatter.py | 2 +- tests/test_instrument_all_and_run.py | 2 +- tests/test_instrument_codeflash_capture.py | 2 +- ...t_instrumentation_run_results_aiservice.py | 2 +- .../test_javascript_optimization_flow.py | 4 +- 33 files changed, 1535 insertions(+), 1485 deletions(-) create mode 100644 codeflash/cli_cmds/github_workflow.py create mode 100644 codeflash/cli_cmds/init_auth.py create mode 100644 codeflash/cli_cmds/init_config.py rename codeflash/{code_utils => cli_cmds}/oauth_handler.py (100%) rename codeflash/{optimization => languages}/function_optimizer.py (100%) rename codeflash/{verification => languages/python}/instrument_codeflash_capture.py (100%) rename codeflash/{verification => languages/python}/parse_line_profile_test_output.py (100%) rename codeflash/{verification => languages/python}/test_runner.py (100%) delete mode 100644 codeflash/optimization/function_context.py diff --git a/.claude/rules/architecture.md b/.claude/rules/architecture.md index 96eb3a43c..1c11d4dc7 100644 --- a/.claude/rules/architecture.md +++ b/.claude/rules/architecture.md @@ -6,10 +6,15 @@ When adding, moving, or deleting source files, update this doc to match. codeflash/ ├── main.py # CLI entry point ├── cli_cmds/ # Command handling, console output (Rich) +│ ├── cmd_init.py # Init orchestrator + Python-specific setup +│ ├── init_config.py # Config types, validation, writing, shared UI helpers +│ ├── init_auth.py # API key management + GitHub app installation +│ ├── github_workflow.py # GitHub Actions workflow generation +│ ├── init_javascript.py # JavaScript/TypeScript project initialization +│ └── oauth_handler.py # OAuth PKCE flow for CodeFlash authentication ├── discovery/ # Find optimizable functions ├── optimization/ # Generate optimized code via AI -│ ├── optimizer.py # Main optimization orchestration -│ └── function_optimizer.py # Per-function optimization logic +│ └── optimizer.py # Main optimization orchestration ├── verification/ # Run deterministic tests (pytest plugin) ├── benchmarking/ # Performance measurement ├── github/ # PR creation @@ -20,12 +25,16 @@ codeflash/ │ ├── base.py # LanguageSupport protocol and shared data types │ ├── registry.py # Language registration and lookup by extension/enum │ ├── current.py # Current language singleton (set_current_language / current_language_support) +│ ├── function_optimizer.py # FunctionOptimizer base class for per-function optimization │ ├── code_replacer.py # Language-agnostic code replacement │ ├── python/ │ │ ├── support.py # PythonSupport (LanguageSupport implementation) │ │ ├── function_optimizer.py # PythonFunctionOptimizer subclass │ │ ├── optimizer.py # Python module preparation & AST resolution -│ │ └── normalizer.py # Python code normalization for deduplication +│ │ ├── normalizer.py # Python code normalization for deduplication +│ │ ├── test_runner.py # Test subprocess execution for Python +│ │ ├── instrument_codeflash_capture.py # Instrument __init__ with capture decorators +│ │ └── parse_line_profile_test_output.py # Parse line profiler output │ └── javascript/ │ ├── support.py # JavaScriptSupport (LanguageSupport implementation) │ ├── function_optimizer.py # JavaScriptFunctionOptimizer subclass @@ -46,9 +55,9 @@ codeflash/ | Task | Start here | |------|------------| -| CLI arguments & commands | `cli_cmds/cli.py` | +| CLI arguments & commands | `cli_cmds/cli.py` (parsing), `main.py` (subcommand dispatch) | | Optimization orchestration | `optimization/optimizer.py` → `run()` | -| Per-function optimization | `optimization/function_optimizer.py` (base), `languages/python/function_optimizer.py`, `languages/javascript/function_optimizer.py` | +| Per-function optimization | `languages/function_optimizer.py` (base), `languages/python/function_optimizer.py`, `languages/javascript/function_optimizer.py` | | Function discovery | `discovery/functions_to_optimize.py` | | Context extraction | `languages//context/code_context_extractor.py` | | Test execution | `languages//support.py` (`run_behavioral_tests`, etc.), `verification/pytest_plugin.py` | diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 0cad9e451..9f5a65251 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -5,10 +5,7 @@ from argparse import SUPPRESS, ArgumentParser, Namespace from pathlib import Path from codeflash.cli_cmds import logging_config -from codeflash.cli_cmds.cli_common import apologize_and_exit -from codeflash.cli_cmds.cmd_init import init_codeflash, install_github_actions from codeflash.cli_cmds.console import logger -from codeflash.cli_cmds.extension import install_vscode_extension from codeflash.code_utils import env_utils from codeflash.code_utils.code_utils import exit_with_message, normalize_ignore_paths from codeflash.code_utils.config_parser import parse_config_file @@ -21,20 +18,12 @@ def parse_args() -> Namespace: parser = ArgumentParser() subparsers = parser.add_subparsers(dest="command", help="Sub-commands") - init_parser = subparsers.add_parser("init", help="Initialize Codeflash for your project.") - init_parser.set_defaults(func=init_codeflash) - + subparsers.add_parser("init", help="Initialize Codeflash for your project.") subparsers.add_parser("vscode-install", help="Install the Codeflash VSCode extension") - - init_actions_parser = subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow") - init_actions_parser.set_defaults(func=install_github_actions) + subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow") trace_optimize = subparsers.add_parser("optimize", help="Trace and optimize your project.") - from codeflash.tracer import main as tracer_main - - trace_optimize.set_defaults(func=tracer_main) - trace_optimize.add_argument( "--max-function-count", type=int, @@ -182,10 +171,6 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace: _handle_reset_config(confirm=not getattr(args, "yes", False)) sys.exit() - if args.command == "vscode-install": - install_vscode_extension() - sys.exit() - if not check_running_in_git_repo(module_root=args.module_root): if not confirm_proceeding_with_no_git_repo(): exit_with_message("No git repository detected and user aborted run. Exiting...", error_on_exit=True) @@ -337,6 +322,8 @@ def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace: f"I couldn't find a git repository in the current directory. " f"I need a git repository to run {mode} and open PRs for optimizations. Exiting..." ) + from codeflash.cli_cmds.cli_common import apologize_and_exit + apologize_and_exit() git_remote = getattr(args, "git_remote", None) if not check_and_push_branch(git_repo, git_remote=git_remote): diff --git a/codeflash/cli_cmds/cmd_init.py b/codeflash/cli_cmds/cmd_init.py index c4e45cc0a..b68807a7b 100644 --- a/codeflash/cli_cmds/cmd_init.py +++ b/codeflash/cli_cmds/cmd_init.py @@ -4,96 +4,44 @@ import os import re import subprocess import sys -from enum import Enum, auto -from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Any, Optional, Union, cast +from typing import TYPE_CHECKING, cast import click -import git import inquirer -import inquirer.themes import tomlkit -from git import InvalidGitRepositoryError, Repo -from pydantic.dataclasses import dataclass from rich.console import Group from rich.panel import Panel from rich.table import Table from rich.text import Text -from codeflash.api.aiservice import AiServiceClient -from codeflash.api.cfapi import get_user_id, is_github_app_installed_on_repo, setup_github_actions from codeflash.cli_cmds.cli_common import apologize_and_exit from codeflash.cli_cmds.console import console, logger from codeflash.cli_cmds.extension import install_vscode_extension - -# Import JS/TS init module -from codeflash.cli_cmds.init_javascript import ( - ProjectLanguage, - detect_project_language, - determine_js_package_manager, - get_js_dependency_installation_commands, - init_js_project, +from codeflash.cli_cmds.github_workflow import install_github_actions +from codeflash.cli_cmds.init_auth import install_github_app, prompt_api_key +from codeflash.cli_cmds.init_config import ( + CLISetupInfo, + CodeflashTheme, + CommonSections, + ask_for_telemetry, + configure_pyproject_toml, + create_empty_pyproject_toml, + get_suggestions, + should_modify_pyproject_toml, ) +from codeflash.cli_cmds.init_javascript import ProjectLanguage, detect_project_language, init_js_project from codeflash.code_utils.code_utils import validate_relative_directory_path from codeflash.code_utils.compat import LF -from codeflash.code_utils.config_parser import parse_config_file -from codeflash.code_utils.env_utils import check_formatter_installed, get_codeflash_api_key -from codeflash.code_utils.git_utils import get_current_branch, get_git_remotes, get_repo_owner_and_name -from codeflash.code_utils.github_utils import get_github_secrets_page_url -from codeflash.code_utils.oauth_handler import perform_oauth_signin -from codeflash.code_utils.shell_utils import get_shell_rc_path, is_powershell, save_api_key_to_rc -from codeflash.either import is_successful +from codeflash.code_utils.env_utils import check_formatter_installed +from codeflash.code_utils.git_utils import get_git_remotes +from codeflash.code_utils.shell_utils import get_shell_rc_path, is_powershell from codeflash.lsp.helpers import is_LSP_enabled from codeflash.telemetry.posthog_cf import ph -from codeflash.version import __version__ as version if TYPE_CHECKING: from argparse import Namespace -CODEFLASH_LOGO: str = ( - f"{LF}" - r" _ ___ _ _ " + f"{LF}" - r" | | / __)| | | | " + f"{LF}" - r" ____ ___ _ | | ____ | |__ | | ____ ___ | | _ " + f"{LF}" - r" / ___) / _ \ / || | / _ )| __)| | / _ | /___)| || \ " + f"{LF}" - r"( (___ | |_| |( (_| |( (/ / | | | |( ( | ||___ || | | |" + f"{LF}" - r" \____) \___/ \____| \____)|_| |_| \_||_|(___/ |_| |_|" + f"{LF}" - f"{('v' + version).rjust(66)}{LF}" - f"{LF}" -) - - -@dataclass(frozen=True) -class CLISetupInfo: - """Setup info for Python projects.""" - - module_root: str - tests_root: str - benchmarks_root: Union[str, None] - ignore_paths: list[str] - formatter: Union[str, list[str]] - git_remote: str - enable_telemetry: bool - - -@dataclass(frozen=True) -class VsCodeSetupInfo: - """Setup info for VSCode extension initialization.""" - - module_root: str - tests_root: str - formatter: Union[str, list[str]] - - -class DependencyManager(Enum): - """Python dependency managers.""" - - PIP = auto() - POETRY = auto() - UV = auto() - UNKNOWN = auto() - def init_codeflash() -> None: try: @@ -110,7 +58,6 @@ def init_codeflash() -> None: console.print(welcome_panel) console.print() - # TODO:{claude} move the init_javascript to the support folder. Move any other language related specific implementation (other than python) to its support. # Detect project language project_language = detect_project_language() @@ -198,129 +145,9 @@ def ask_run_end_to_end_test(args: Namespace) -> None: run_end_to_end_test(args, file_path) -def config_found(pyproject_toml_path: Union[str, Path]) -> tuple[bool, str]: - pyproject_toml_path = Path(pyproject_toml_path) - - if not pyproject_toml_path.exists(): - return False, f"Configuration file not found: {pyproject_toml_path}" - - if not pyproject_toml_path.is_file(): - return False, f"Configuration file is not a file: {pyproject_toml_path}" - - if pyproject_toml_path.suffix != ".toml": - return False, f"Configuration file is not a .toml file: {pyproject_toml_path}" - - return True, "" - - -def is_valid_pyproject_toml(pyproject_toml_path: Union[str, Path]) -> tuple[bool, dict[str, Any] | None, str]: - pyproject_toml_path = Path(pyproject_toml_path) - try: - config, _ = parse_config_file(pyproject_toml_path) - except Exception as e: - return False, None, f"Failed to parse configuration: {e}" - - module_root = config.get("module_root") - if not module_root: - return False, config, "Missing required field: 'module_root'" - - if not Path(module_root).is_dir(): - return False, config, f"Invalid 'module_root': directory does not exist at {module_root}" - - tests_root = config.get("tests_root") - if not tests_root: - return False, config, "Missing required field: 'tests_root'" - - if not Path(tests_root).is_dir(): - return False, config, f"Invalid 'tests_root': directory does not exist at {tests_root}" - - return True, 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. - - If it does, ask the user if they want to re-configure it. - """ - from rich.prompt import Confirm - - pyproject_toml_path = Path.cwd() / "pyproject.toml" - - found, _ = config_found(pyproject_toml_path) - if not found: - return True, None - - valid, config, _message = is_valid_pyproject_toml(pyproject_toml_path) - if not valid: - # needs to be re-configured - return True, None - - return Confirm.ask( - "✅ A valid Codeflash config already exists in this project. Do you want to re-configure it?", - default=False, - show_default=True, - ), config - - -# Custom theme for better UX -class CodeflashTheme(inquirer.themes.Default): - def __init__(self) -> None: - super().__init__() - self.Question.mark_color = inquirer.themes.term.yellow - self.Question.brackets_color = inquirer.themes.term.bright_blue - self.Question.default_color = inquirer.themes.term.bright_cyan - self.List.selection_color = inquirer.themes.term.bright_blue - self.Checkbox.selection_color = inquirer.themes.term.bright_blue - self.Checkbox.selected_icon = "✅" - self.Checkbox.unselected_icon = "⬜" - - -# common sections between normal mode and lsp mode -class CommonSections(Enum): - module_root = "module_root" - tests_root = "tests_root" - formatter_cmds = "formatter_cmds" - - def get_toml_key(self) -> str: - return self.value.replace("_", "-") - - -@lru_cache(maxsize=1) -def get_valid_subdirs(current_dir: Optional[Path] = None) -> list[str]: - ignore_subdirs = [ - "venv", - "node_modules", - "dist", - "build", - "build_temp", - "build_scripts", - "env", - "logs", - "tmp", - "__pycache__", - ] - path_str = str(current_dir) if current_dir else "." - return [ - d - for d in next(os.walk(path_str))[1] - if not d.startswith(".") and not d.startswith("__") and d not in ignore_subdirs - ] - - -def get_suggestions(section: str) -> tuple[list[str], Optional[str]]: - valid_subdirs = get_valid_subdirs() - if section == CommonSections.module_root: - return [d for d in valid_subdirs if d != "tests"], None - if section == CommonSections.tests_root: - default = "tests" if "tests" in valid_subdirs else None - return valid_subdirs, default - if section == CommonSections.formatter_cmds: - return ["disabled", "ruff", "black"], "disabled" - msg = f"Unknown section: {section}" - raise ValueError(msg) - - def collect_setup_info() -> CLISetupInfo: + from git import InvalidGitRepositoryError, Repo + curdir = Path.cwd() # Check if the cwd is writable if not os.access(curdir, os.W_OK): @@ -500,20 +327,6 @@ def collect_setup_info() -> CLISetupInfo: benchmarks_root = None - # TODO: Implement other benchmark framework options - # if benchmarks_root: - # benchmarks_root = benchmarks_root.relative_to(curdir) - # - # # Ask about benchmark framework - # benchmark_framework_options = ["pytest-benchmark", "asv (Airspeed Velocity)", "custom/other"] - # benchmark_framework = inquirer_wrapper( - # inquirer.list_input, - # message="Which benchmark framework do you use?", - # choices=benchmark_framework_options, - # default=benchmark_framework_options[0], - # carousel=True, - # ) - formatter_panel = Panel( Text( "🎨 Let's configure your code formatter.\n\n" @@ -699,1175 +512,6 @@ def check_for_toml_or_setup_file() -> str | None: return cast("str", project_name) -def create_empty_pyproject_toml(pyproject_toml_path: Path) -> None: - ph("cli-create-pyproject-toml") - lsp_mode = is_LSP_enabled() - # Define a minimal pyproject.toml content - new_pyproject_toml = tomlkit.document() - new_pyproject_toml["tool"] = {"codeflash": {}} - try: - pyproject_toml_path.write_text(tomlkit.dumps(new_pyproject_toml), encoding="utf8") - - # Check if the pyproject.toml file was created - if pyproject_toml_path.exists() and not lsp_mode: - success_panel = Panel( - Text( - f"✅ Created a pyproject.toml file at {pyproject_toml_path}\n\n" - "Your project is now ready for Codeflash configuration!", - style="green", - justify="center", - ), - title="🎉 Success!", - border_style="bright_green", - ) - console.print(success_panel) - console.print("\n📍 Press any key to continue...") - console.input() - ph("cli-created-pyproject-toml") - except OSError: - click.echo("❌ Failed to create pyproject.toml. Please check your disk permissions and available space.") - apologize_and_exit() - - -def install_github_actions(override_formatter_check: bool = False) -> None: - try: - config, _config_file_path = parse_config_file(override_formatter_check=override_formatter_check) - - ph("cli-github-actions-install-started") - try: - repo = Repo(config["module_root"], search_parent_directories=True) - except git.InvalidGitRepositoryError: - click.echo( - "Skipping GitHub action installation for continuous optimization because you're not in a git repository." - ) - return - - git_root = Path(repo.git.rev_parse("--show-toplevel")) - workflows_path = git_root / ".github" / "workflows" - optimize_yaml_path = workflows_path / "codeflash.yaml" - - # Check if workflow file already exists locally BEFORE showing prompt - if optimize_yaml_path.exists(): - # Workflow file already exists locally - skip prompt and setup - already_exists_message = "✅ GitHub Actions workflow file already exists.\n\n" - already_exists_message += "No changes needed - your repository is already configured!" - - already_exists_panel = Panel( - Text(already_exists_message, style="green", justify="center"), - title="✅ Already Configured", - border_style="bright_green", - ) - console.print(already_exists_panel) - console.print() - - logger.info("[cmd_init.py:install_github_actions] Workflow file already exists locally, skipping setup") - return - - # Get repository information for API call - git_remote = config.get("git_remote", "origin") - # get_current_branch handles detached HEAD and other edge cases internally - try: - base_branch = get_current_branch(repo) - except Exception as e: - logger.warning( - f"[cmd_init.py:install_github_actions] Could not determine current branch: {e}. Falling back to 'main'." - ) - base_branch = "main" - - # Generate workflow content - from importlib.resources import files - - benchmark_mode = False - benchmarks_root = config.get("benchmarks_root", "").strip() - if benchmarks_root and benchmarks_root != "": - benchmark_panel = Panel( - Text( - "📊 Benchmark Mode Available\n\n" - "I noticed you've configured a benchmarks_root in your config. " - "Benchmark mode will show the performance impact of Codeflash's optimizations on your benchmarks.", - style="cyan", - ), - title="📊 Benchmark Mode", - border_style="bright_cyan", - ) - console.print(benchmark_panel) - console.print() - - benchmark_questions = [ - inquirer.Confirm("benchmark_mode", message="Run GitHub Actions in benchmark mode?", default=True) - ] - - benchmark_answers = inquirer.prompt(benchmark_questions, theme=CodeflashTheme()) - benchmark_mode = benchmark_answers["benchmark_mode"] if benchmark_answers else False - - # Show prompt only if workflow doesn't exist locally - actions_panel = Panel( - Text( - "🤖 GitHub Actions Setup\n\n" - "GitHub Actions will automatically optimize your code in every pull request. " - "This is the recommended way to use Codeflash for continuous optimization.", - style="blue", - ), - title="🤖 Continuous Optimization", - border_style="bright_blue", - ) - console.print(actions_panel) - console.print() - - creation_questions = [ - inquirer.Confirm( - "confirm_creation", - message="Set up GitHub Actions for continuous optimization? We'll open a pull request with the workflow file.", - default=True, - ) - ] - - creation_answers = inquirer.prompt(creation_questions, theme=CodeflashTheme()) - if not creation_answers or not creation_answers["confirm_creation"]: - skip_panel = Panel( - Text("⏩️ Skipping GitHub Actions setup.", style="yellow"), title="⏩️ Skipped", border_style="yellow" - ) - console.print(skip_panel) - ph("cli-github-workflow-skipped") - return - ph( - "cli-github-optimization-confirm-workflow-creation", - {"confirm_creation": creation_answers["confirm_creation"]}, - ) - - # Generate workflow content AFTER user confirmation - logger.info("[cmd_init.py:install_github_actions] User confirmed, generating workflow content...") - - # Select the appropriate workflow template based on project language - project_language = detect_project_language_for_workflow(Path.cwd()) - if project_language in ("javascript", "typescript"): - workflow_template = "codeflash-optimize-js.yaml" - else: - workflow_template = "codeflash-optimize.yaml" - - optimize_yml_content = ( - files("codeflash").joinpath("cli_cmds", "workflows", workflow_template).read_text(encoding="utf-8") - ) - materialized_optimize_yml_content = generate_dynamic_workflow_content( - optimize_yml_content, config, git_root, benchmark_mode - ) - - workflows_path.mkdir(parents=True, exist_ok=True) - - pr_created_via_api = False - pr_url = None - - try: - owner, repo_name = get_repo_owner_and_name(repo, git_remote) - except Exception as e: - logger.error(f"[cmd_init.py:install_github_actions] Failed to get repository owner and name: {e}") - # Fall back to local file creation - workflows_path.mkdir(parents=True, exist_ok=True) - with optimize_yaml_path.open("w", encoding="utf8") as optimize_yml_file: - optimize_yml_file.write(materialized_optimize_yml_content) - workflow_success_panel = Panel( - Text( - f"✅ Created GitHub action workflow at {optimize_yaml_path}\n\n" - "Your repository is now configured for continuous optimization!", - style="green", - justify="center", - ), - title="🎉 Workflow Created!", - border_style="bright_green", - ) - console.print(workflow_success_panel) - console.print() - else: - # Try to create PR via API - try: - # Workflow file doesn't exist on remote or content differs - proceed with PR creation - console.print("Creating PR with GitHub Actions workflow...") - logger.info( - f"[cmd_init.py:install_github_actions] Calling setup_github_actions API for {owner}/{repo_name} on branch {base_branch}" - ) - - response = setup_github_actions( - owner=owner, - repo=repo_name, - base_branch=base_branch, - workflow_content=materialized_optimize_yml_content, - ) - - if response.status_code == 200: - response_data = response.json() - if response_data.get("success"): - pr_url = response_data.get("pr_url") - - if pr_url: - pr_created_via_api = True - success_message = f"✅ PR created: {pr_url}\n\n" - success_message += "Your repository is now configured for continuous optimization!" - - workflow_success_panel = Panel( - Text(success_message, style="green", justify="center"), - title="🎉 Workflow PR Created!", - border_style="bright_green", - ) - console.print(workflow_success_panel) - console.print() - - logger.info( - f"[cmd_init.py:install_github_actions] Successfully created PR #{response_data.get('pr_number')} for {owner}/{repo_name}" - ) - else: - # File already exists with same content - pr_created_via_api = True # Mark as handled (no PR needed) - already_exists_message = "✅ Workflow file already exists with the same content.\n\n" - already_exists_message += "No changes needed - your repository is already configured!" - - already_exists_panel = Panel( - Text(already_exists_message, style="green", justify="center"), - title="✅ Already Configured", - border_style="bright_green", - ) - console.print(already_exists_panel) - console.print() - else: - # API returned success=false, extract error details - error_data = response_data - error_msg = error_data.get("error", "Unknown error") - error_message = error_data.get("message", error_msg) - error_help = error_data.get("help", "") - installation_url = error_data.get("installation_url") - - # For permission errors, don't fall back - show a focused message and abort early - if response.status_code == 403: - logger.error( - f"[cmd_init.py:install_github_actions] Permission denied for {owner}/{repo_name}" - ) - # Extract installation_url if available, otherwise use default - installation_url_403 = error_data.get( - "installation_url", "https://github.com/apps/codeflash-ai/installations/select_target" - ) - - permission_error_panel = Panel( - Text( - "❌ Access Denied\n\n" - f"The GitHub App may not be installed on {owner}/{repo_name}, or it doesn't have the required permissions.\n\n" - "💡 To fix this:\n" - "1. Install the CodeFlash GitHub App on your repository\n" - "2. Ensure the app has 'Contents: write', 'Workflows: write', and 'Pull requests: write' permissions\n" - "3. Make sure you have write access to the repository\n\n" - f"🔗 Install GitHub App: {installation_url_403}", - style="red", - ), - title="❌ Setup Failed", - border_style="red", - ) - console.print(permission_error_panel) - console.print() - click.echo( - f"Please install the CodeFlash GitHub App and ensure it has the required permissions.{LF}" - f"Visit: {installation_url_403}{LF}" - ) - apologize_and_exit() - - # Show detailed error panel for all other errors - error_panel_text = f"❌ {error_msg}\n\n{error_message}\n" - if error_help: - error_panel_text += f"\n💡 {error_help}\n" - if installation_url: - error_panel_text += f"\n🔗 Install GitHub App: {installation_url}" - - error_panel = Panel( - Text(error_panel_text, style="red"), title="❌ Setup Failed", border_style="red" - ) - console.print(error_panel) - console.print() - - # For GitHub App not installed, don't fall back - show clear instructions - if response.status_code == 404 and installation_url: - logger.error( - f"[cmd_init.py:install_github_actions] GitHub App not installed on {owner}/{repo_name}" - ) - click.echo( - f"Please install the CodeFlash GitHub App on your repository to continue.{LF}" - f"Visit: {installation_url}{LF}" - ) - return - - # For other errors, fall back to local file creation - raise Exception(error_message) # noqa: TRY002, TRY301 - else: - # API call returned non-200 status, try to parse error response - try: - error_data = response.json() - error_msg = error_data.get("error", "API request failed") - error_message = error_data.get("message", f"API returned status {response.status_code}") - error_help = error_data.get("help", "") - installation_url = error_data.get("installation_url") - - # For permission errors, don't fall back - show a focused message and abort early - if response.status_code == 403: - logger.error( - f"[cmd_init.py:install_github_actions] Permission denied for {owner}/{repo_name}" - ) - # Extract installation_url if available, otherwise use default - installation_url_403 = error_data.get( - "installation_url", "https://github.com/apps/codeflash-ai/installations/select_target" - ) - - permission_error_panel = Panel( - Text( - "❌ Access Denied\n\n" - f"The GitHub App may not be installed on {owner}/{repo_name}, or it doesn't have the required permissions.\n\n" - "💡 To fix this:\n" - "1. Install the CodeFlash GitHub App on your repository\n" - "2. Ensure the app has 'Contents: write', 'Workflows: write', and 'Pull requests: write' permissions\n" - "3. Make sure you have write access to the repository\n\n" - f"🔗 Install GitHub App: {installation_url_403}", - style="red", - ), - title="❌ Setup Failed", - border_style="red", - ) - console.print(permission_error_panel) - console.print() - click.echo( - f"Please install the CodeFlash GitHub App and ensure it has the required permissions.{LF}" - f"Visit: {installation_url_403}{LF}" - ) - apologize_and_exit() - - # Show detailed error panel for all other errors - error_panel_text = f"❌ {error_msg}\n\n{error_message}\n" - if error_help: - error_panel_text += f"\n💡 {error_help}\n" - if installation_url: - error_panel_text += f"\n🔗 Install GitHub App: {installation_url}" - - error_panel = Panel( - Text(error_panel_text, style="red"), title="❌ Setup Failed", border_style="red" - ) - console.print(error_panel) - console.print() - - # For GitHub App not installed, don't fall back - show clear instructions - if response.status_code == 404 and installation_url: - logger.error( - f"[cmd_init.py:install_github_actions] GitHub App not installed on {owner}/{repo_name}" - ) - click.echo( - f"Please install the CodeFlash GitHub App on your repository to continue.{LF}" - f"Visit: {installation_url}{LF}" - ) - return - - # For authentication errors, don't fall back - if response.status_code == 401: - logger.error( - f"[cmd_init.py:install_github_actions] Authentication failed for {owner}/{repo_name}" - ) - click.echo(f"Authentication failed. Please check your API key and try again.{LF}") - return - - # For other errors, fall back to local file creation - raise Exception(error_message) # noqa: TRY002 - except (ValueError, KeyError) as parse_error: - # Couldn't parse error response, use generic message - status_msg = f"API returned status {response.status_code}" - raise Exception(status_msg) from parse_error # noqa: TRY002 - - except Exception as api_error: - # Fall back to local file creation if API call fails (for non-critical errors) - logger.warning( - f"[cmd_init.py:install_github_actions] API call failed, falling back to local file creation: {api_error}" - ) - workflows_path.mkdir(parents=True, exist_ok=True) - with optimize_yaml_path.open("w", encoding="utf8") as optimize_yml_file: - optimize_yml_file.write(materialized_optimize_yml_content) - workflow_success_panel = Panel( - Text( - f"✅ Created GitHub action workflow at {optimize_yaml_path}\n\n" - "Your repository is now configured for continuous optimization!", - style="green", - justify="center", - ), - title="🎉 Workflow Created!", - border_style="bright_green", - ) - console.print(workflow_success_panel) - console.print() - - # Show appropriate message based on whether PR was created via API - if pr_created_via_api: - if pr_url: - click.echo( - f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}" - f"Once you merge the PR, the workflow will be active.{LF}" - ) - else: - # File already exists - click.echo( - f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}" - f"The workflow is ready to use.{LF}" - ) - else: - # Fell back to local file creation - click.echo( - f"Please edit, commit and push this GitHub actions file to your repo, and you're all set!{LF}" - f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}" - ) - - # Show GitHub secrets setup panel (needed in both cases - PR created via API or local file) - try: - existing_api_key = get_codeflash_api_key() - except OSError: - existing_api_key = None - - # GitHub secrets setup panel - always shown since secrets are required for the workflow to work - secrets_message = ( - "🔐 Next Step: Add API Key as GitHub Secret\n\n" - "You'll need to add your CODEFLASH_API_KEY as a secret to your GitHub repository.\n\n" - "📋 Steps:\n" - "1. Press Enter to open your repo's secrets page\n" - "2. Click 'New repository secret'\n" - "3. Add your API key with the variable name CODEFLASH_API_KEY" - ) - - if existing_api_key: - secrets_message += f"\n\n🔑 Your API Key: {existing_api_key}" - - secrets_panel = Panel( - Text(secrets_message, style="blue"), title="🔐 GitHub Secrets Setup", border_style="bright_blue" - ) - console.print(secrets_panel) - - console.print(f"\n📍 Press Enter to open: {get_github_secrets_page_url(repo)}") - console.input() - - click.launch(get_github_secrets_page_url(repo)) - - # Post-launch message panel - launch_panel = Panel( - Text( - "🐙 I opened your GitHub secrets page!\n\n" - "Note: If you see a 404, you probably don't have access to this repo's secrets. " - "Ask a repo admin to add it for you, or (not recommended) you can temporarily " - "hard-code your API key into the workflow file.", - style="cyan", - ), - title="🌐 Browser Opened", - border_style="bright_cyan", - ) - console.print(launch_panel) - click.pause() - console.print() - ph("cli-github-workflow-created") - except KeyboardInterrupt: - apologize_and_exit() - - -def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager: - """Determine which dependency manager is being used based on pyproject.toml contents.""" - cwd = Path.cwd() - if (cwd / "poetry.lock").exists(): - return DependencyManager.POETRY - if (cwd / "uv.lock").exists(): - return DependencyManager.UV - if "tool" not in pyproject_data: - return DependencyManager.PIP - - tool_section = pyproject_data["tool"] - - # Check for poetry - if "poetry" in tool_section: - return DependencyManager.POETRY - - # Check for uv - if any(key.startswith("uv") for key in tool_section): - return DependencyManager.UV - - # Look for pip-specific markers - if "pip" in tool_section or "setuptools" in tool_section: - return DependencyManager.PIP - - return DependencyManager.UNKNOWN - - -def get_codeflash_github_action_command(dep_manager: DependencyManager) -> str: - """Generate the appropriate codeflash command based on the dependency manager.""" - if dep_manager == DependencyManager.POETRY: - return """| - poetry env use python - poetry run codeflash""" - if dep_manager == DependencyManager.UV: - return "uv run codeflash" - # PIP or UNKNOWN - return "codeflash" - - -def get_dependency_installation_commands(dep_manager: DependencyManager) -> tuple[str, str]: - """Generate commands to install the dependency manager and project dependencies.""" - if dep_manager == DependencyManager.POETRY: - return """| - python -m pip install --upgrade pip - pip install poetry - poetry install --all-extras""" - if dep_manager == DependencyManager.UV: - return """| - uv sync --all-extras - uv pip install --upgrade codeflash""" - # PIP or UNKNOWN - return """| - python -m pip install --upgrade pip - pip install -r requirements.txt - pip install codeflash""" - - -def get_dependency_manager_installation_string(dep_manager: DependencyManager) -> str: - py_version = sys.version_info - python_version_string = f"'{py_version.major}.{py_version.minor}'" - if dep_manager == DependencyManager.UV: - return """name: 🐍 Setup UV - uses: astral-sh/setup-uv@v6 - with: - enable-cache: true""" - return f"""name: 🐍 Set up Python - uses: actions/setup-python@v5 - with: - python-version: {python_version_string}""" - - -def get_github_action_working_directory(toml_path: Path, git_root: Path) -> str: - if toml_path.parent == git_root: - return "" - working_dir = str(toml_path.parent.relative_to(git_root)) - return f"""defaults: - run: - working-directory: ./{working_dir}""" - - -# ============================================================================ -# JavaScript/TypeScript GitHub Actions Support -# ============================================================================ -# Note: JS package manager and workflow helper functions are imported from init_javascript.py - - -def detect_project_language_for_workflow(project_root: Path) -> str: - """Detect the primary language of the project for workflow generation. - - Returns: 'python', 'javascript', or 'typescript' - """ - # Check for TypeScript config - if (project_root / "tsconfig.json").exists(): - return "typescript" - - # Check for JavaScript/TypeScript indicators - has_package_json = (project_root / "package.json").exists() - has_pyproject = (project_root / "pyproject.toml").exists() - - if has_package_json and not has_pyproject: - # Pure JS/TS project - return "javascript" - if has_pyproject and not has_package_json: - # Pure Python project - return "python" - - # Both exist - count files to determine primary language - js_count = 0 - py_count = 0 - for file in project_root.rglob("*"): - if file.is_file(): - suffix = file.suffix.lower() - if suffix in {".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs"}: - js_count += 1 - elif suffix == ".py": - py_count += 1 - - if js_count > py_count: - return "javascript" - return "python" - - -def collect_repo_files_for_workflow(git_root: Path) -> dict[str, Any]: - """Collect important repository files and directory structure for workflow generation. - - :param git_root: Root directory of the git repository - :return: Dictionary with 'files' (path -> content) and 'directory_structure' (nested dict) - """ - # Important files to collect with contents - important_files = [ - "pyproject.toml", - "requirements.txt", - "requirements-dev.txt", - "requirements/requirements.txt", - "requirements/dev.txt", - "Pipfile", - "Pipfile.lock", - "poetry.lock", - "uv.lock", - "setup.py", - "setup.cfg", - "Dockerfile", - "docker-compose.yml", - "docker-compose.yaml", - "Makefile", - "README.md", - "README.rst", - ] - - # Also collect GitHub workflows - workflows_path = git_root / ".github" / "workflows" - if workflows_path.exists(): - important_files.extend( - str(workflow_file.relative_to(git_root)) for workflow_file in workflows_path.glob("*.yml") - ) - important_files.extend( - str(workflow_file.relative_to(git_root)) for workflow_file in workflows_path.glob("*.yaml") - ) - - files_dict: dict[str, str] = {} - max_file_size = 8 * 1024 # 8KB limit per file - - for file_path_str in important_files: - file_path = git_root / file_path_str - if file_path.exists() and file_path.is_file(): - try: - content = file_path.read_text(encoding="utf-8", errors="ignore") - # Limit file size - if len(content) > max_file_size: - content = content[:max_file_size] + "\n... (truncated)" - files_dict[file_path_str] = content - except Exception as e: - logger.warning(f"[cmd_init.py:collect_repo_files_for_workflow] Failed to read {file_path_str}: {e}") - - # Collect 2-level directory structure - directory_structure: dict[str, Any] = {} - try: - for item in sorted(git_root.iterdir()): - if item.name.startswith(".") and item.name not in [".github", ".git"]: - continue # Skip hidden files/folders except .github - - if item.is_dir(): - # Level 1: directory - dir_dict: dict[str, Any] = {"type": "directory", "contents": {}} - try: - # Level 2: contents of directory - for subitem in sorted(item.iterdir()): - if subitem.name.startswith("."): - continue - if subitem.is_dir(): - dir_dict["contents"][subitem.name] = {"type": "directory"} - else: - dir_dict["contents"][subitem.name] = {"type": "file"} - except PermissionError: - pass # Skip directories we can't read - directory_structure[item.name] = dir_dict - elif item.is_file(): - directory_structure[item.name] = {"type": "file"} - except Exception as e: - logger.warning(f"[cmd_init.py:collect_repo_files_for_workflow] Error collecting directory structure: {e}") - - return {"files": files_dict, "directory_structure": directory_structure} - - -def generate_dynamic_workflow_content( - optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False -) -> str: - """Generate workflow content with dynamic steps from AI service, falling back to static template. - - :param optimize_yml_content: Base workflow template content - :param config: Codeflash configuration tuple (dict, Path) - :param git_root: Root directory of the git repository - :param benchmark_mode: Whether to enable benchmark mode - :return: Complete workflow YAML content - """ - # First, do the basic replacements that are always needed - module_path = str(Path(config["module_root"]).relative_to(git_root) / "**") - optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path) - - # Detect project language - project_language = detect_project_language_for_workflow(Path.cwd()) - - # For JavaScript/TypeScript projects, use static template customization - # (AI-generated steps are currently Python-only) - if project_language in ("javascript", "typescript"): - return customize_codeflash_yaml_content(optimize_yml_content, config, git_root, benchmark_mode) - - # Python project - try AI-generated steps - toml_path = Path.cwd() / "pyproject.toml" - try: - with toml_path.open(encoding="utf8") as pyproject_file: - pyproject_data = tomlkit.parse(pyproject_file.read()) - except FileNotFoundError: - click.echo( - f"I couldn't find a pyproject.toml in the current directory.{LF}" - f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file." - ) - apologize_and_exit() - - working_dir = get_github_action_working_directory(toml_path, git_root) - optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir) - - # Try to generate dynamic steps using AI service - try: - repo_data = collect_repo_files_for_workflow(git_root) - - # Prepare codeflash config for AI - codeflash_config = { - "module_root": config["module_root"], - "tests_root": config.get("tests_root", ""), - "benchmark_mode": benchmark_mode, - } - - aiservice_client = AiServiceClient() - dynamic_steps = aiservice_client.generate_workflow_steps( - repo_files=repo_data["files"], - directory_structure=repo_data["directory_structure"], - codeflash_config=codeflash_config, - ) - - if dynamic_steps: - # Replace the entire steps section with AI-generated steps - # Find the steps section in the template - steps_start = optimize_yml_content.find(" steps:") - if steps_start != -1: - # Find the end of the steps section (next line at same or less indentation) - lines = optimize_yml_content.split("\n") - steps_start_line = optimize_yml_content[:steps_start].count("\n") - steps_end_line = len(lines) - - # Find where steps section ends (next job or end of file) - for i in range(steps_start_line + 1, len(lines)): - line = lines[i] - # Stop if we hit a line that's not indented (new job or end of jobs) - if line and not line.startswith(" ") and not line.startswith("\t"): - steps_end_line = i - break - - # Extract steps content from AI response (remove "steps:" prefix if present) - steps_content = dynamic_steps - if steps_content.startswith("steps:"): - # Remove "steps:" and leading newline - steps_content = steps_content[6:].lstrip("\n") - - # Ensure proper indentation (8 spaces for steps section in YAML) - indented_steps = [] - for line in steps_content.split("\n"): - if line.strip(): - # If line doesn't start with enough spaces, add them - if not line.startswith(" "): - indented_steps.append(" " + line) - else: - # Preserve existing indentation but ensure minimum 8 spaces - current_indent = len(line) - len(line.lstrip()) - if current_indent < 8: - indented_steps.append(" " * 8 + line.lstrip()) - else: - indented_steps.append(line) - else: - indented_steps.append("") - - # Add codeflash command step at the end - dep_manager = determine_dependency_manager(pyproject_data) - codeflash_cmd = get_codeflash_github_action_command(dep_manager) - if benchmark_mode: - codeflash_cmd += " --benchmark" - - # Format codeflash command properly - if "|" in codeflash_cmd: - # Multi-line command - cmd_lines = codeflash_cmd.split("\n") - codeflash_step = f" - name: ⚡️Codeflash Optimization\n run: {cmd_lines[0].strip()}" - for cmd_line in cmd_lines[1:]: - codeflash_step += f"\n {cmd_line.strip()}" - else: - codeflash_step = f" - name: ⚡️Codeflash Optimization\n run: {codeflash_cmd}" - - indented_steps.append(codeflash_step) - - # Reconstruct the workflow - return "\n".join([*lines[:steps_start_line], " steps:", *indented_steps, *lines[steps_end_line:]]) - logger.warning("[cmd_init.py:generate_dynamic_workflow_content] Could not find steps section in template") - else: - logger.debug( - "[cmd_init.py:generate_dynamic_workflow_content] AI service returned no steps, falling back to static" - ) - - except Exception as e: - logger.warning( - f"[cmd_init.py:generate_dynamic_workflow_content] Error generating dynamic workflow, falling back to static: {e}" - ) - - # Fallback to static template - return customize_codeflash_yaml_content(optimize_yml_content, config, git_root, benchmark_mode) - - -def customize_codeflash_yaml_content( - optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False -) -> str: - module_path = str(Path(config["module_root"]).relative_to(git_root) / "**") - optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path) - - # Detect project language - project_language = detect_project_language_for_workflow(Path.cwd()) - - if project_language in ("javascript", "typescript"): - # JavaScript/TypeScript project - return _customize_js_workflow_content(optimize_yml_content, git_root, benchmark_mode) - - # Python project (default) - return _customize_python_workflow_content(optimize_yml_content, git_root, benchmark_mode) - - -def _customize_python_workflow_content(optimize_yml_content: str, git_root: Path, benchmark_mode: bool = False) -> str: - """Customize workflow content for Python projects.""" - # Get dependency installation commands - toml_path = Path.cwd() / "pyproject.toml" - try: - with toml_path.open(encoding="utf8") as pyproject_file: - pyproject_data = tomlkit.parse(pyproject_file.read()) - except FileNotFoundError: - click.echo( - f"I couldn't find a pyproject.toml in the current directory.{LF}" - f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file." - ) - apologize_and_exit() - - working_dir = get_github_action_working_directory(toml_path, git_root) - optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir) - dep_manager = determine_dependency_manager(pyproject_data) - - python_depmanager_installation = get_dependency_manager_installation_string(dep_manager) - optimize_yml_content = optimize_yml_content.replace( - "{{ setup_runtime_environment }}", python_depmanager_installation - ) - install_deps_cmd = get_dependency_installation_commands(dep_manager) - - optimize_yml_content = optimize_yml_content.replace("{{ install_dependencies_command }}", install_deps_cmd) - - # Add codeflash command - codeflash_cmd = get_codeflash_github_action_command(dep_manager) - - if benchmark_mode: - codeflash_cmd += " --benchmark" - return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd) - - -# TODO:{claude} Refactor and move to support for language specific -def _customize_js_workflow_content(optimize_yml_content: str, git_root: Path, benchmark_mode: bool = False) -> str: - """Customize workflow content for JavaScript/TypeScript projects.""" - from codeflash.cli_cmds.init_javascript import ( - get_js_codeflash_install_step, - get_js_codeflash_run_command, - get_js_runtime_setup_steps, - is_codeflash_dependency, - ) - - project_root = Path.cwd() - package_json_path = project_root / "package.json" - - if not package_json_path.exists(): - click.echo( - f"I couldn't find a package.json in the current directory.{LF}" - f"Please run `npm init` or create a package.json file first." - ) - apologize_and_exit() - - # Determine working directory relative to git root - if project_root == git_root: - working_dir = "" - else: - rel_path = str(project_root.relative_to(git_root)) - working_dir = f"""defaults: - run: - working-directory: ./{rel_path}""" - - optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir) - - # Determine package manager and codeflash dependency status - pkg_manager = determine_js_package_manager(project_root) - codeflash_is_dep = is_codeflash_dependency(project_root) - - # Setup runtime environment (Node.js/Bun) - runtime_setup = get_js_runtime_setup_steps(pkg_manager) - optimize_yml_content = optimize_yml_content.replace("{{ setup_runtime_steps }}", runtime_setup) - - # Install dependencies - install_deps_cmd = get_js_dependency_installation_commands(pkg_manager) - optimize_yml_content = optimize_yml_content.replace("{{ install_dependencies_command }}", install_deps_cmd) - - # Install codeflash step (only if not a dependency) - install_codeflash = get_js_codeflash_install_step(pkg_manager, is_dependency=codeflash_is_dep) - optimize_yml_content = optimize_yml_content.replace("{{ install_codeflash_step }}", install_codeflash) - - # Codeflash run command - codeflash_cmd = get_js_codeflash_run_command(pkg_manager, is_dependency=codeflash_is_dep) - if benchmark_mode: - codeflash_cmd += " --benchmark" - return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd) - - -def get_formatter_cmds(formatter: str) -> list[str]: - if formatter == "black": - return ["black $file"] - if formatter == "ruff": - return ["ruff check --exit-zero --fix $file", "ruff format $file"] - if formatter == "other": - click.echo( - "🔧 In pyproject.toml, please replace 'your-formatter' with the command you use to format your code." - ) - return ["your-formatter $file"] - if formatter in {"don't use a formatter", "disabled"}: - return ["disabled"] - if " && " in formatter: - return formatter.split(" && ") - return [formatter] - - -# Create or update the pyproject.toml file with the Codeflash dependency & configuration -def configure_pyproject_toml( - setup_info: Union[VsCodeSetupInfo, CLISetupInfo], config_file: Optional[Path] = None -) -> bool: - for_vscode = isinstance(setup_info, VsCodeSetupInfo) - toml_path = config_file or Path.cwd() / "pyproject.toml" - try: - with toml_path.open(encoding="utf8") as pyproject_file: - pyproject_data = tomlkit.parse(pyproject_file.read()) - except FileNotFoundError: - click.echo( - f"I couldn't find a pyproject.toml in the current directory.{LF}" - f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file." - ) - return False - - codeflash_section = tomlkit.table() - codeflash_section.add(tomlkit.comment("All paths are relative to this pyproject.toml's directory.")) - - if for_vscode: - for section in CommonSections: - if hasattr(setup_info, section.value): - codeflash_section[section.get_toml_key()] = getattr(setup_info, section.value) - else: - codeflash_section["module-root"] = setup_info.module_root - codeflash_section["tests-root"] = setup_info.tests_root - codeflash_section["ignore-paths"] = setup_info.ignore_paths - if not setup_info.enable_telemetry: - codeflash_section["disable-telemetry"] = not setup_info.enable_telemetry - if setup_info.git_remote not in ["", "origin"]: - codeflash_section["git-remote"] = setup_info.git_remote - - formatter = setup_info.formatter - - formatter_cmds = formatter if isinstance(formatter, list) else get_formatter_cmds(formatter) - - check_formatter_installed(formatter_cmds, exit_on_failure=False) - codeflash_section["formatter-cmds"] = formatter_cmds - # Add the 'codeflash' section, ensuring 'tool' section exists - tool_section = pyproject_data.get("tool", tomlkit.table()) - - if for_vscode: - # merge the existing codeflash section, instead of overwriting it - existing_codeflash = tool_section.get("codeflash", tomlkit.table()) - - for key, value in codeflash_section.items(): - existing_codeflash[key] = value - tool_section["codeflash"] = existing_codeflash - else: - tool_section["codeflash"] = codeflash_section - - pyproject_data["tool"] = tool_section - - with toml_path.open("w", encoding="utf8") as pyproject_file: - pyproject_file.write(tomlkit.dumps(pyproject_data)) - click.echo(f"Added Codeflash configuration to {toml_path}") - click.echo() - return True - - -def install_github_app(git_remote: str) -> None: - try: - git_repo = git.Repo(search_parent_directories=True) - 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): - click.echo( - f"🐙 Looks like you've already installed the Codeflash GitHub app on this repository ({owner}/{repo})! Continuing…" - ) - - else: - try: - click.prompt( - f"Finally, you'll need to install the Codeflash GitHub app by choosing the repository you want to install Codeflash on.{LF}" - f"I will attempt to open the github app page - https://github.com/apps/codeflash-ai/installations/select_target {LF}" - f"Please, press ENTER to open the app installation page{LF}", - default="", - type=click.STRING, - prompt_suffix=">>> ", - show_default=False, - ) - click.launch("https://github.com/apps/codeflash-ai/installations/select_target") - click.prompt( - f"Please, press ENTER once you've finished installing the github app from https://github.com/apps/codeflash-ai/installations/select_target{LF}", - default="", - type=click.STRING, - prompt_suffix=">>> ", - show_default=False, - ) - - count = 2 - while not is_github_app_installed_on_repo(owner, repo, suppress_errors=True): - if count == 0: - click.echo( - f"❌ It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo}.{LF}" - f"You won't be able to create PRs with Codeflash until you install the app.{LF}" - f"In the meantime you can make local only optimizations by using the '--no-pr' flag with codeflash.{LF}" - ) - break - click.prompt( - f"❌ It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo}.{LF}" - f"Please install it from https://github.com/apps/codeflash-ai/installations/select_target {LF}" - f"Please, press ENTER to continue once you've finished installing the github app…{LF}", - default="", - type=click.STRING, - prompt_suffix=">>> ", - show_default=False, - ) - count -= 1 - except (KeyboardInterrupt, EOFError, click.exceptions.Abort): - # leave empty line for the next prompt to be properly rendered - click.echo() - - -class CFAPIKeyType(click.ParamType): - name = "cfapi-key" - - def convert(self, value: str, param: click.Parameter | None, ctx: click.Context | None) -> str | None: - value = value.strip() - if not value.startswith("cf-") and value != "": - self.fail( - f"That key [{value}] seems to be invalid. It should start with a 'cf-' prefix. Please try again.", - param, - ctx, - ) - return value - - -# Returns True if the user entered a new API key, False if they used an existing one -def prompt_api_key() -> bool: - """Prompt user for API key via OAuth or manual entry.""" - # Check for existing API key - try: - existing_api_key = get_codeflash_api_key() - except OSError: - existing_api_key = None - - if existing_api_key: - display_key = f"{existing_api_key[:3]}****{existing_api_key[-4:]}" - api_key_panel = Panel( - Text( - f"🔑 I found a CODEFLASH_API_KEY in your environment [{display_key}]!\n\n" - "✅ You're all set with API authentication!", - style="green", - justify="center", - ), - title="🔑 API Key Found", - border_style="bright_green", - ) - console.print(api_key_panel) - console.print() - return False - - # Prompt for authentication method - auth_choices = ["🔐 Login in with Codeflash", "🔑 Use Codeflash API key"] - - questions = [ - inquirer.List( - "auth_method", - message="How would you like to authenticate?", - choices=auth_choices, - default=auth_choices[0], - carousel=True, - ) - ] - - answers = inquirer.prompt(questions, theme=CodeflashTheme()) - if not answers: - apologize_and_exit() - - method = answers["auth_method"] - - if method == auth_choices[1]: - enter_api_key_and_save_to_rc() - ph("cli-new-api-key-entered") - return True - - # Perform OAuth sign-in - api_key = perform_oauth_signin() - - if not api_key: - apologize_and_exit() - - # Save API key - shell_rc_path = get_shell_rc_path() - if not shell_rc_path.exists() and os.name == "nt": - shell_rc_path.touch() - click.echo(f"✅ Created {shell_rc_path}") - - result = save_api_key_to_rc(api_key) - if is_successful(result): - click.echo(result.unwrap()) - click.echo("✅ Signed in successfully and API key saved!") - else: - click.echo(result.failure()) - click.pause() - - os.environ["CODEFLASH_API_KEY"] = api_key - ph("cli-oauth-signin-completed") - return True - - -def enter_api_key_and_save_to_rc() -> None: - browser_launched = False - api_key = "" - while api_key == "": - api_key = click.prompt( - f"Enter your Codeflash API key{' [or press Enter to open your API key page]' if not browser_launched else ''}", - hide_input=False, - default="", - type=CFAPIKeyType(), - show_default=False, - ).strip() - if api_key: - break - if not browser_launched: - click.echo( - f"Opening your Codeflash API key page. Grab a key from there!{LF}" - "You can also open this link manually: https://app.codeflash.ai/app/apikeys" - ) - click.launch("https://app.codeflash.ai/app/apikeys") - browser_launched = True # This does not work on remote consoles - shell_rc_path = get_shell_rc_path() - if not shell_rc_path.exists() and os.name == "nt": - # On Windows, create the appropriate file (PowerShell .ps1 or CMD .bat) in the user's home directory - shell_rc_path.parent.mkdir(parents=True, exist_ok=True) - shell_rc_path.touch() - click.echo(f"✅ Created {shell_rc_path}") - get_user_id(api_key=api_key) # Used to verify whether the API key is valid. - result = save_api_key_to_rc(api_key) - if is_successful(result): - click.echo(result.unwrap()) - else: - click.echo(result.failure()) - click.pause() - - os.environ["CODEFLASH_API_KEY"] = api_key - - def create_find_common_tags_file(args: Namespace, file_name: str) -> Path: find_common_tags_content = """from __future__ import annotations @@ -1993,14 +637,3 @@ def run_end_to_end_test(args: Namespace, find_common_tags_path: Path) -> None: logger.info("🧹 Cleaning up…") find_common_tags_path.unlink(missing_ok=True) logger.info(f"🗑️ Deleted {find_common_tags_path}") - - -def ask_for_telemetry() -> bool: - """Prompt the user to enable or disable telemetry.""" - from rich.prompt import Confirm - - return Confirm.ask( - "⚡️ Help us improve Codeflash by sharing anonymous usage data (e.g. errors encountered)?", - default=True, - show_default=True, - ) diff --git a/codeflash/cli_cmds/console_constants.py b/codeflash/cli_cmds/console_constants.py index 13f8aa621..7363108e9 100644 --- a/codeflash/cli_cmds/console_constants.py +++ b/codeflash/cli_cmds/console_constants.py @@ -1,3 +1,18 @@ +from codeflash.code_utils.compat import LF +from codeflash.version import __version__ as version + +CODEFLASH_LOGO: str = ( + f"{LF}" + r" _ ___ _ _ " + f"{LF}" + r" | | / __)| | | | " + f"{LF}" + r" ____ ___ _ | | ____ | |__ | | ____ ___ | | _ " + f"{LF}" + r" / ___) / _ \ / || | / _ )| __)| | / _ | /___)| || \ " + f"{LF}" + r"( (___ | |_| |( (_| |( (/ / | | | |( ( | ||___ || | | |" + f"{LF}" + r" \____) \___/ \____| \____)|_| |_| \_||_|(___/ |_| |_|" + f"{LF}" + f"{('v' + version).rjust(66)}{LF}" + f"{LF}" +) + SPINNER_TYPES = { "point", "simpleDots", diff --git a/codeflash/cli_cmds/github_workflow.py b/codeflash/cli_cmds/github_workflow.py new file mode 100644 index 000000000..5abb095dd --- /dev/null +++ b/codeflash/cli_cmds/github_workflow.py @@ -0,0 +1,898 @@ +from __future__ import annotations + +import sys +from enum import Enum, auto +from pathlib import Path +from typing import Any + +import click +import git +import inquirer +import tomlkit +from git import Repo +from rich.panel import Panel +from rich.text import Text + +from codeflash.api.aiservice import AiServiceClient +from codeflash.api.cfapi import setup_github_actions +from codeflash.cli_cmds.cli_common import apologize_and_exit +from codeflash.cli_cmds.console import console, logger +from codeflash.cli_cmds.init_config import CodeflashTheme +from codeflash.code_utils.compat import LF +from codeflash.code_utils.config_parser import parse_config_file +from codeflash.code_utils.env_utils import get_codeflash_api_key +from codeflash.code_utils.git_utils import get_current_branch, get_repo_owner_and_name +from codeflash.code_utils.github_utils import get_github_secrets_page_url +from codeflash.telemetry.posthog_cf import ph + + +class DependencyManager(Enum): + """Python dependency managers.""" + + PIP = auto() + POETRY = auto() + UV = auto() + UNKNOWN = auto() + + +def install_github_actions(override_formatter_check: bool = False) -> None: + try: + config, _config_file_path = parse_config_file(override_formatter_check=override_formatter_check) + + ph("cli-github-actions-install-started") + try: + repo = Repo(config["module_root"], search_parent_directories=True) + except git.InvalidGitRepositoryError: + click.echo( + "Skipping GitHub action installation for continuous optimization because you're not in a git repository." + ) + return + + git_root = Path(repo.git.rev_parse("--show-toplevel")) + workflows_path = git_root / ".github" / "workflows" + optimize_yaml_path = workflows_path / "codeflash.yaml" + + # Check if workflow file already exists locally BEFORE showing prompt + if optimize_yaml_path.exists(): + # Workflow file already exists locally - skip prompt and setup + already_exists_message = "✅ GitHub Actions workflow file already exists.\n\n" + already_exists_message += "No changes needed - your repository is already configured!" + + already_exists_panel = Panel( + Text(already_exists_message, style="green", justify="center"), + title="✅ Already Configured", + border_style="bright_green", + ) + console.print(already_exists_panel) + console.print() + + logger.info("[github_workflow.py:install_github_actions] Workflow file already exists locally, skipping setup") + return + + # Get repository information for API call + git_remote = config.get("git_remote", "origin") + # get_current_branch handles detached HEAD and other edge cases internally + try: + base_branch = get_current_branch(repo) + except Exception as e: + logger.warning( + f"[github_workflow.py:install_github_actions] Could not determine current branch: {e}. Falling back to 'main'." + ) + base_branch = "main" + + # Generate workflow content + from importlib.resources import files + + benchmark_mode = False + benchmarks_root = config.get("benchmarks_root", "").strip() + if benchmarks_root and benchmarks_root != "": + benchmark_panel = Panel( + Text( + "📊 Benchmark Mode Available\n\n" + "I noticed you've configured a benchmarks_root in your config. " + "Benchmark mode will show the performance impact of Codeflash's optimizations on your benchmarks.", + style="cyan", + ), + title="📊 Benchmark Mode", + border_style="bright_cyan", + ) + console.print(benchmark_panel) + console.print() + + benchmark_questions = [ + inquirer.Confirm("benchmark_mode", message="Run GitHub Actions in benchmark mode?", default=True) + ] + + benchmark_answers = inquirer.prompt(benchmark_questions, theme=CodeflashTheme()) + benchmark_mode = benchmark_answers["benchmark_mode"] if benchmark_answers else False + + # Show prompt only if workflow doesn't exist locally + actions_panel = Panel( + Text( + "🤖 GitHub Actions Setup\n\n" + "GitHub Actions will automatically optimize your code in every pull request. " + "This is the recommended way to use Codeflash for continuous optimization.", + style="blue", + ), + title="🤖 Continuous Optimization", + border_style="bright_blue", + ) + console.print(actions_panel) + console.print() + + creation_questions = [ + inquirer.Confirm( + "confirm_creation", + message="Set up GitHub Actions for continuous optimization? We'll open a pull request with the workflow file.", + default=True, + ) + ] + + creation_answers = inquirer.prompt(creation_questions, theme=CodeflashTheme()) + if not creation_answers or not creation_answers["confirm_creation"]: + skip_panel = Panel( + Text("⏩️ Skipping GitHub Actions setup.", style="yellow"), title="⏩️ Skipped", border_style="yellow" + ) + console.print(skip_panel) + ph("cli-github-workflow-skipped") + return + ph( + "cli-github-optimization-confirm-workflow-creation", + {"confirm_creation": creation_answers["confirm_creation"]}, + ) + + # Generate workflow content AFTER user confirmation + logger.info("[github_workflow.py:install_github_actions] User confirmed, generating workflow content...") + + # Select the appropriate workflow template based on project language + project_language = detect_project_language_for_workflow(Path.cwd()) + if project_language in ("javascript", "typescript"): + workflow_template = "codeflash-optimize-js.yaml" + else: + workflow_template = "codeflash-optimize.yaml" + + optimize_yml_content = ( + files("codeflash").joinpath("cli_cmds", "workflows", workflow_template).read_text(encoding="utf-8") + ) + materialized_optimize_yml_content = generate_dynamic_workflow_content( + optimize_yml_content, config, git_root, benchmark_mode + ) + + workflows_path.mkdir(parents=True, exist_ok=True) + + pr_created_via_api = False + pr_url = None + + try: + owner, repo_name = get_repo_owner_and_name(repo, git_remote) + except Exception as e: + logger.error(f"[github_workflow.py:install_github_actions] Failed to get repository owner and name: {e}") + # Fall back to local file creation + workflows_path.mkdir(parents=True, exist_ok=True) + with optimize_yaml_path.open("w", encoding="utf8") as optimize_yml_file: + optimize_yml_file.write(materialized_optimize_yml_content) + workflow_success_panel = Panel( + Text( + f"✅ Created GitHub action workflow at {optimize_yaml_path}\n\n" + "Your repository is now configured for continuous optimization!", + style="green", + justify="center", + ), + title="🎉 Workflow Created!", + border_style="bright_green", + ) + console.print(workflow_success_panel) + console.print() + else: + # Try to create PR via API + try: + # Workflow file doesn't exist on remote or content differs - proceed with PR creation + console.print("Creating PR with GitHub Actions workflow...") + logger.info( + f"[github_workflow.py:install_github_actions] Calling setup_github_actions API for {owner}/{repo_name} on branch {base_branch}" + ) + + response = setup_github_actions( + owner=owner, + repo=repo_name, + base_branch=base_branch, + workflow_content=materialized_optimize_yml_content, + ) + + if response.status_code == 200: + response_data = response.json() + if response_data.get("success"): + pr_url = response_data.get("pr_url") + + if pr_url: + pr_created_via_api = True + success_message = f"✅ PR created: {pr_url}\n\n" + success_message += "Your repository is now configured for continuous optimization!" + + workflow_success_panel = Panel( + Text(success_message, style="green", justify="center"), + title="🎉 Workflow PR Created!", + border_style="bright_green", + ) + console.print(workflow_success_panel) + console.print() + + logger.info( + f"[github_workflow.py:install_github_actions] Successfully created PR #{response_data.get('pr_number')} for {owner}/{repo_name}" + ) + else: + # File already exists with same content + pr_created_via_api = True # Mark as handled (no PR needed) + already_exists_message = "✅ Workflow file already exists with the same content.\n\n" + already_exists_message += "No changes needed - your repository is already configured!" + + already_exists_panel = Panel( + Text(already_exists_message, style="green", justify="center"), + title="✅ Already Configured", + border_style="bright_green", + ) + console.print(already_exists_panel) + console.print() + else: + # API returned success=false, extract error details + error_data = response_data + error_msg = error_data.get("error", "Unknown error") + error_message = error_data.get("message", error_msg) + error_help = error_data.get("help", "") + installation_url = error_data.get("installation_url") + + # For permission errors, don't fall back - show a focused message and abort early + if response.status_code == 403: + logger.error( + f"[github_workflow.py:install_github_actions] Permission denied for {owner}/{repo_name}" + ) + # Extract installation_url if available, otherwise use default + installation_url_403 = error_data.get( + "installation_url", "https://github.com/apps/codeflash-ai/installations/select_target" + ) + + permission_error_panel = Panel( + Text( + "❌ Access Denied\n\n" + f"The GitHub App may not be installed on {owner}/{repo_name}, or it doesn't have the required permissions.\n\n" + "💡 To fix this:\n" + "1. Install the CodeFlash GitHub App on your repository\n" + "2. Ensure the app has 'Contents: write', 'Workflows: write', and 'Pull requests: write' permissions\n" + "3. Make sure you have write access to the repository\n\n" + f"🔗 Install GitHub App: {installation_url_403}", + style="red", + ), + title="❌ Setup Failed", + border_style="red", + ) + console.print(permission_error_panel) + console.print() + click.echo( + f"Please install the CodeFlash GitHub App and ensure it has the required permissions.{LF}" + f"Visit: {installation_url_403}{LF}" + ) + apologize_and_exit() + + # Show detailed error panel for all other errors + error_panel_text = f"❌ {error_msg}\n\n{error_message}\n" + if error_help: + error_panel_text += f"\n💡 {error_help}\n" + if installation_url: + error_panel_text += f"\n🔗 Install GitHub App: {installation_url}" + + error_panel = Panel( + Text(error_panel_text, style="red"), title="❌ Setup Failed", border_style="red" + ) + console.print(error_panel) + console.print() + + # For GitHub App not installed, don't fall back - show clear instructions + if response.status_code == 404 and installation_url: + logger.error( + f"[github_workflow.py:install_github_actions] GitHub App not installed on {owner}/{repo_name}" + ) + click.echo( + f"Please install the CodeFlash GitHub App on your repository to continue.{LF}" + f"Visit: {installation_url}{LF}" + ) + return + + # For other errors, fall back to local file creation + raise Exception(error_message) # noqa: TRY002, TRY301 + else: + # API call returned non-200 status, try to parse error response + try: + error_data = response.json() + error_msg = error_data.get("error", "API request failed") + error_message = error_data.get("message", f"API returned status {response.status_code}") + error_help = error_data.get("help", "") + installation_url = error_data.get("installation_url") + + # For permission errors, don't fall back - show a focused message and abort early + if response.status_code == 403: + logger.error( + f"[github_workflow.py:install_github_actions] Permission denied for {owner}/{repo_name}" + ) + # Extract installation_url if available, otherwise use default + installation_url_403 = error_data.get( + "installation_url", "https://github.com/apps/codeflash-ai/installations/select_target" + ) + + permission_error_panel = Panel( + Text( + "❌ Access Denied\n\n" + f"The GitHub App may not be installed on {owner}/{repo_name}, or it doesn't have the required permissions.\n\n" + "💡 To fix this:\n" + "1. Install the CodeFlash GitHub App on your repository\n" + "2. Ensure the app has 'Contents: write', 'Workflows: write', and 'Pull requests: write' permissions\n" + "3. Make sure you have write access to the repository\n\n" + f"🔗 Install GitHub App: {installation_url_403}", + style="red", + ), + title="❌ Setup Failed", + border_style="red", + ) + console.print(permission_error_panel) + console.print() + click.echo( + f"Please install the CodeFlash GitHub App and ensure it has the required permissions.{LF}" + f"Visit: {installation_url_403}{LF}" + ) + apologize_and_exit() + + # Show detailed error panel for all other errors + error_panel_text = f"❌ {error_msg}\n\n{error_message}\n" + if error_help: + error_panel_text += f"\n💡 {error_help}\n" + if installation_url: + error_panel_text += f"\n🔗 Install GitHub App: {installation_url}" + + error_panel = Panel( + Text(error_panel_text, style="red"), title="❌ Setup Failed", border_style="red" + ) + console.print(error_panel) + console.print() + + # For GitHub App not installed, don't fall back - show clear instructions + if response.status_code == 404 and installation_url: + logger.error( + f"[github_workflow.py:install_github_actions] GitHub App not installed on {owner}/{repo_name}" + ) + click.echo( + f"Please install the CodeFlash GitHub App on your repository to continue.{LF}" + f"Visit: {installation_url}{LF}" + ) + return + + # For authentication errors, don't fall back + if response.status_code == 401: + logger.error( + f"[github_workflow.py:install_github_actions] Authentication failed for {owner}/{repo_name}" + ) + click.echo(f"Authentication failed. Please check your API key and try again.{LF}") + return + + # For other errors, fall back to local file creation + raise Exception(error_message) # noqa: TRY002 + except (ValueError, KeyError) as parse_error: + # Couldn't parse error response, use generic message + status_msg = f"API returned status {response.status_code}" + raise Exception(status_msg) from parse_error # noqa: TRY002 + + except Exception as api_error: + # Fall back to local file creation if API call fails (for non-critical errors) + logger.warning( + f"[github_workflow.py:install_github_actions] API call failed, falling back to local file creation: {api_error}" + ) + workflows_path.mkdir(parents=True, exist_ok=True) + with optimize_yaml_path.open("w", encoding="utf8") as optimize_yml_file: + optimize_yml_file.write(materialized_optimize_yml_content) + workflow_success_panel = Panel( + Text( + f"✅ Created GitHub action workflow at {optimize_yaml_path}\n\n" + "Your repository is now configured for continuous optimization!", + style="green", + justify="center", + ), + title="🎉 Workflow Created!", + border_style="bright_green", + ) + console.print(workflow_success_panel) + console.print() + + # Show appropriate message based on whether PR was created via API + if pr_created_via_api: + if pr_url: + click.echo( + f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}" + f"Once you merge the PR, the workflow will be active.{LF}" + ) + else: + # File already exists + click.echo( + f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}" + f"The workflow is ready to use.{LF}" + ) + else: + # Fell back to local file creation + click.echo( + f"Please edit, commit and push this GitHub actions file to your repo, and you're all set!{LF}" + f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}" + ) + + # Show GitHub secrets setup panel (needed in both cases - PR created via API or local file) + try: + existing_api_key = get_codeflash_api_key() + except OSError: + existing_api_key = None + + # GitHub secrets setup panel - always shown since secrets are required for the workflow to work + secrets_message = ( + "🔐 Next Step: Add API Key as GitHub Secret\n\n" + "You'll need to add your CODEFLASH_API_KEY as a secret to your GitHub repository.\n\n" + "📋 Steps:\n" + "1. Press Enter to open your repo's secrets page\n" + "2. Click 'New repository secret'\n" + "3. Add your API key with the variable name CODEFLASH_API_KEY" + ) + + if existing_api_key: + secrets_message += f"\n\n🔑 Your API Key: {existing_api_key}" + + secrets_panel = Panel( + Text(secrets_message, style="blue"), title="🔐 GitHub Secrets Setup", border_style="bright_blue" + ) + console.print(secrets_panel) + + console.print(f"\n📍 Press Enter to open: {get_github_secrets_page_url(repo)}") + console.input() + + click.launch(get_github_secrets_page_url(repo)) + + # Post-launch message panel + launch_panel = Panel( + Text( + "🐙 I opened your GitHub secrets page!\n\n" + "Note: If you see a 404, you probably don't have access to this repo's secrets. " + "Ask a repo admin to add it for you, or (not recommended) you can temporarily " + "hard-code your API key into the workflow file.", + style="cyan", + ), + title="🌐 Browser Opened", + border_style="bright_cyan", + ) + console.print(launch_panel) + click.pause() + console.print() + ph("cli-github-workflow-created") + except KeyboardInterrupt: + apologize_and_exit() + + +def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager: + """Determine which dependency manager is being used based on pyproject.toml contents.""" + cwd = Path.cwd() + if (cwd / "poetry.lock").exists(): + return DependencyManager.POETRY + if (cwd / "uv.lock").exists(): + return DependencyManager.UV + if "tool" not in pyproject_data: + return DependencyManager.PIP + + tool_section = pyproject_data["tool"] + + # Check for poetry + if "poetry" in tool_section: + return DependencyManager.POETRY + + # Check for uv + if any(key.startswith("uv") for key in tool_section): + return DependencyManager.UV + + # Look for pip-specific markers + if "pip" in tool_section or "setuptools" in tool_section: + return DependencyManager.PIP + + return DependencyManager.UNKNOWN + + +def get_codeflash_github_action_command(dep_manager: DependencyManager) -> str: + """Generate the appropriate codeflash command based on the dependency manager.""" + if dep_manager == DependencyManager.POETRY: + return """| + poetry env use python + poetry run codeflash""" + if dep_manager == DependencyManager.UV: + return "uv run codeflash" + # PIP or UNKNOWN + return "codeflash" + + +def get_dependency_installation_commands(dep_manager: DependencyManager) -> tuple[str, str]: + """Generate commands to install the dependency manager and project dependencies.""" + if dep_manager == DependencyManager.POETRY: + return """| + python -m pip install --upgrade pip + pip install poetry + poetry install --all-extras""" + if dep_manager == DependencyManager.UV: + return """| + uv sync --all-extras + uv pip install --upgrade codeflash""" + # PIP or UNKNOWN + return """| + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install codeflash""" + + +def get_dependency_manager_installation_string(dep_manager: DependencyManager) -> str: + py_version = sys.version_info + python_version_string = f"'{py_version.major}.{py_version.minor}'" + if dep_manager == DependencyManager.UV: + return """name: 🐍 Setup UV + uses: astral-sh/setup-uv@v6 + with: + enable-cache: true""" + return f"""name: 🐍 Set up Python + uses: actions/setup-python@v5 + with: + python-version: {python_version_string}""" + + +def get_github_action_working_directory(toml_path: Path, git_root: Path) -> str: + if toml_path.parent == git_root: + return "" + working_dir = str(toml_path.parent.relative_to(git_root)) + return f"""defaults: + run: + working-directory: ./{working_dir}""" + + +def detect_project_language_for_workflow(project_root: Path) -> str: + """Detect the primary language of the project for workflow generation. + + Returns: 'python', 'javascript', or 'typescript' + """ + # Check for TypeScript config + if (project_root / "tsconfig.json").exists(): + return "typescript" + + # Check for JavaScript/TypeScript indicators + has_package_json = (project_root / "package.json").exists() + has_pyproject = (project_root / "pyproject.toml").exists() + + if has_package_json and not has_pyproject: + # Pure JS/TS project + return "javascript" + if has_pyproject and not has_package_json: + # Pure Python project + return "python" + + # Both exist - count files to determine primary language + js_count = 0 + py_count = 0 + for file in project_root.rglob("*"): + if file.is_file(): + suffix = file.suffix.lower() + if suffix in {".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs"}: + js_count += 1 + elif suffix == ".py": + py_count += 1 + + if js_count > py_count: + return "javascript" + return "python" + + +def collect_repo_files_for_workflow(git_root: Path) -> dict[str, Any]: + """Collect important repository files and directory structure for workflow generation. + + :param git_root: Root directory of the git repository + :return: Dictionary with 'files' (path -> content) and 'directory_structure' (nested dict) + """ + # Important files to collect with contents + important_files = [ + "pyproject.toml", + "requirements.txt", + "requirements-dev.txt", + "requirements/requirements.txt", + "requirements/dev.txt", + "Pipfile", + "Pipfile.lock", + "poetry.lock", + "uv.lock", + "setup.py", + "setup.cfg", + "Dockerfile", + "docker-compose.yml", + "docker-compose.yaml", + "Makefile", + "README.md", + "README.rst", + ] + + # Also collect GitHub workflows + workflows_path = git_root / ".github" / "workflows" + if workflows_path.exists(): + important_files.extend( + str(workflow_file.relative_to(git_root)) for workflow_file in workflows_path.glob("*.yml") + ) + important_files.extend( + str(workflow_file.relative_to(git_root)) for workflow_file in workflows_path.glob("*.yaml") + ) + + files_dict: dict[str, str] = {} + max_file_size = 8 * 1024 # 8KB limit per file + + for file_path_str in important_files: + file_path = git_root / file_path_str + if file_path.exists() and file_path.is_file(): + try: + content = file_path.read_text(encoding="utf-8", errors="ignore") + # Limit file size + if len(content) > max_file_size: + content = content[:max_file_size] + "\n... (truncated)" + files_dict[file_path_str] = content + except Exception as e: + logger.warning(f"[github_workflow.py:collect_repo_files_for_workflow] Failed to read {file_path_str}: {e}") + + # Collect 2-level directory structure + directory_structure: dict[str, Any] = {} + try: + for item in sorted(git_root.iterdir()): + if item.name.startswith(".") and item.name not in [".github", ".git"]: + continue # Skip hidden files/folders except .github + + if item.is_dir(): + # Level 1: directory + dir_dict: dict[str, Any] = {"type": "directory", "contents": {}} + try: + # Level 2: contents of directory + for subitem in sorted(item.iterdir()): + if subitem.name.startswith("."): + continue + if subitem.is_dir(): + dir_dict["contents"][subitem.name] = {"type": "directory"} + else: + dir_dict["contents"][subitem.name] = {"type": "file"} + except PermissionError: + pass # Skip directories we can't read + directory_structure[item.name] = dir_dict + elif item.is_file(): + directory_structure[item.name] = {"type": "file"} + except Exception as e: + logger.warning(f"[github_workflow.py:collect_repo_files_for_workflow] Error collecting directory structure: {e}") + + return {"files": files_dict, "directory_structure": directory_structure} + + +def generate_dynamic_workflow_content( + optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False +) -> str: + """Generate workflow content with dynamic steps from AI service, falling back to static template.""" + # First, do the basic replacements that are always needed + module_path = str(Path(config["module_root"]).relative_to(git_root) / "**") + optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path) + + # Detect project language + project_language = detect_project_language_for_workflow(Path.cwd()) + + # For JavaScript/TypeScript projects, use static template customization + # (AI-generated steps are currently Python-only) + if project_language in ("javascript", "typescript"): + return customize_codeflash_yaml_content(optimize_yml_content, config, git_root, benchmark_mode) + + # Python project - try AI-generated steps + toml_path = Path.cwd() / "pyproject.toml" + try: + with toml_path.open(encoding="utf8") as pyproject_file: + pyproject_data = tomlkit.parse(pyproject_file.read()) + except FileNotFoundError: + click.echo( + f"I couldn't find a pyproject.toml in the current directory.{LF}" + f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file." + ) + apologize_and_exit() + + working_dir = get_github_action_working_directory(toml_path, git_root) + optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir) + + # Try to generate dynamic steps using AI service + try: + repo_data = collect_repo_files_for_workflow(git_root) + + # Prepare codeflash config for AI + codeflash_config = { + "module_root": config["module_root"], + "tests_root": config.get("tests_root", ""), + "benchmark_mode": benchmark_mode, + } + + aiservice_client = AiServiceClient() + dynamic_steps = aiservice_client.generate_workflow_steps( + repo_files=repo_data["files"], + directory_structure=repo_data["directory_structure"], + codeflash_config=codeflash_config, + ) + + if dynamic_steps: + # Replace the entire steps section with AI-generated steps + # Find the steps section in the template + steps_start = optimize_yml_content.find(" steps:") + if steps_start != -1: + # Find the end of the steps section (next line at same or less indentation) + lines = optimize_yml_content.split("\n") + steps_start_line = optimize_yml_content[:steps_start].count("\n") + steps_end_line = len(lines) + + # Find where steps section ends (next job or end of file) + for i in range(steps_start_line + 1, len(lines)): + line = lines[i] + # Stop if we hit a line that's not indented (new job or end of jobs) + if line and not line.startswith(" ") and not line.startswith("\t"): + steps_end_line = i + break + + # Extract steps content from AI response (remove "steps:" prefix if present) + steps_content = dynamic_steps + if steps_content.startswith("steps:"): + # Remove "steps:" and leading newline + steps_content = steps_content[6:].lstrip("\n") + + # Ensure proper indentation (8 spaces for steps section in YAML) + indented_steps = [] + for line in steps_content.split("\n"): + if line.strip(): + # If line doesn't start with enough spaces, add them + if not line.startswith(" "): + indented_steps.append(" " + line) + else: + # Preserve existing indentation but ensure minimum 8 spaces + current_indent = len(line) - len(line.lstrip()) + if current_indent < 8: + indented_steps.append(" " * 8 + line.lstrip()) + else: + indented_steps.append(line) + else: + indented_steps.append("") + + # Add codeflash command step at the end + dep_manager = determine_dependency_manager(pyproject_data) + codeflash_cmd = get_codeflash_github_action_command(dep_manager) + if benchmark_mode: + codeflash_cmd += " --benchmark" + + # Format codeflash command properly + if "|" in codeflash_cmd: + # Multi-line command + cmd_lines = codeflash_cmd.split("\n") + codeflash_step = f" - name: ⚡️Codeflash Optimization\n run: {cmd_lines[0].strip()}" + for cmd_line in cmd_lines[1:]: + codeflash_step += f"\n {cmd_line.strip()}" + else: + codeflash_step = f" - name: ⚡️Codeflash Optimization\n run: {codeflash_cmd}" + + indented_steps.append(codeflash_step) + + # Reconstruct the workflow + return "\n".join([*lines[:steps_start_line], " steps:", *indented_steps, *lines[steps_end_line:]]) + logger.warning("[github_workflow.py:generate_dynamic_workflow_content] Could not find steps section in template") + else: + logger.debug( + "[github_workflow.py:generate_dynamic_workflow_content] AI service returned no steps, falling back to static" + ) + + except Exception as e: + logger.warning( + f"[github_workflow.py:generate_dynamic_workflow_content] Error generating dynamic workflow, falling back to static: {e}" + ) + + # Fallback to static template + return customize_codeflash_yaml_content(optimize_yml_content, config, git_root, benchmark_mode) + + +def customize_codeflash_yaml_content( + optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False +) -> str: + module_path = str(Path(config["module_root"]).relative_to(git_root) / "**") + optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path) + + # Detect project language + project_language = detect_project_language_for_workflow(Path.cwd()) + + if project_language in ("javascript", "typescript"): + # JavaScript/TypeScript project + return _customize_js_workflow_content(optimize_yml_content, git_root, benchmark_mode) + + # Python project (default) + return _customize_python_workflow_content(optimize_yml_content, git_root, benchmark_mode) + + +def _customize_python_workflow_content(optimize_yml_content: str, git_root: Path, benchmark_mode: bool = False) -> str: + """Customize workflow content for Python projects.""" + # Get dependency installation commands + toml_path = Path.cwd() / "pyproject.toml" + try: + with toml_path.open(encoding="utf8") as pyproject_file: + pyproject_data = tomlkit.parse(pyproject_file.read()) + except FileNotFoundError: + click.echo( + f"I couldn't find a pyproject.toml in the current directory.{LF}" + f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file." + ) + apologize_and_exit() + + working_dir = get_github_action_working_directory(toml_path, git_root) + optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir) + dep_manager = determine_dependency_manager(pyproject_data) + + python_depmanager_installation = get_dependency_manager_installation_string(dep_manager) + optimize_yml_content = optimize_yml_content.replace( + "{{ setup_runtime_environment }}", python_depmanager_installation + ) + install_deps_cmd = get_dependency_installation_commands(dep_manager) + + optimize_yml_content = optimize_yml_content.replace("{{ install_dependencies_command }}", install_deps_cmd) + + # Add codeflash command + codeflash_cmd = get_codeflash_github_action_command(dep_manager) + + if benchmark_mode: + codeflash_cmd += " --benchmark" + return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd) + + +def _customize_js_workflow_content(optimize_yml_content: str, git_root: Path, benchmark_mode: bool = False) -> str: + """Customize workflow content for JavaScript/TypeScript projects.""" + from codeflash.cli_cmds.init_javascript import ( + determine_js_package_manager, + get_js_codeflash_install_step, + get_js_codeflash_run_command, + get_js_dependency_installation_commands, + get_js_runtime_setup_steps, + is_codeflash_dependency, + ) + + project_root = Path.cwd() + package_json_path = project_root / "package.json" + + if not package_json_path.exists(): + click.echo( + f"I couldn't find a package.json in the current directory.{LF}" + f"Please run `npm init` or create a package.json file first." + ) + apologize_and_exit() + + # Determine working directory relative to git root + if project_root == git_root: + working_dir = "" + else: + rel_path = str(project_root.relative_to(git_root)) + working_dir = f"""defaults: + run: + working-directory: ./{rel_path}""" + + optimize_yml_content = optimize_yml_content.replace("{{ working_directory }}", working_dir) + + # Determine package manager and codeflash dependency status + pkg_manager = determine_js_package_manager(project_root) + codeflash_is_dep = is_codeflash_dependency(project_root) + + # Setup runtime environment (Node.js/Bun) + runtime_setup = get_js_runtime_setup_steps(pkg_manager) + optimize_yml_content = optimize_yml_content.replace("{{ setup_runtime_steps }}", runtime_setup) + + # Install dependencies + install_deps_cmd = get_js_dependency_installation_commands(pkg_manager) + optimize_yml_content = optimize_yml_content.replace("{{ install_dependencies_command }}", install_deps_cmd) + + # Install codeflash step (only if not a dependency) + install_codeflash = get_js_codeflash_install_step(pkg_manager, is_dependency=codeflash_is_dep) + optimize_yml_content = optimize_yml_content.replace("{{ install_codeflash_step }}", install_codeflash) + + # Codeflash run command + codeflash_cmd = get_js_codeflash_run_command(pkg_manager, is_dependency=codeflash_is_dep) + if benchmark_mode: + codeflash_cmd += " --benchmark" + return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd) diff --git a/codeflash/cli_cmds/init_auth.py b/codeflash/cli_cmds/init_auth.py new file mode 100644 index 000000000..6dfdf470f --- /dev/null +++ b/codeflash/cli_cmds/init_auth.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import os + +import click +import git +import inquirer + +from codeflash.api.cfapi import get_user_id, is_github_app_installed_on_repo +from codeflash.cli_cmds.cli_common import apologize_and_exit +from codeflash.cli_cmds.console import console +from codeflash.cli_cmds.init_config import CodeflashTheme +from codeflash.cli_cmds.oauth_handler import perform_oauth_signin +from codeflash.code_utils.compat import LF +from codeflash.code_utils.env_utils import get_codeflash_api_key +from codeflash.code_utils.git_utils import get_git_remotes, get_repo_owner_and_name +from codeflash.code_utils.shell_utils import get_shell_rc_path, save_api_key_to_rc +from codeflash.either import is_successful +from codeflash.telemetry.posthog_cf import ph + + +class CFAPIKeyType(click.ParamType): + name = "cfapi-key" + + def convert(self, value: str, param: click.Parameter | None, ctx: click.Context | None) -> str | None: + value = value.strip() + if not value.startswith("cf-") and value != "": + self.fail( + f"That key [{value}] seems to be invalid. It should start with a 'cf-' prefix. Please try again.", + param, + ctx, + ) + return value + + +# Returns True if the user entered a new API key, False if they used an existing one +def prompt_api_key() -> bool: + """Prompt user for API key via OAuth or manual entry.""" + from rich.panel import Panel + from rich.text import Text + + # Check for existing API key + try: + existing_api_key = get_codeflash_api_key() + except OSError: + existing_api_key = None + + if existing_api_key: + display_key = f"{existing_api_key[:3]}****{existing_api_key[-4:]}" + api_key_panel = Panel( + Text( + f"🔑 I found a CODEFLASH_API_KEY in your environment [{display_key}]!\n\n" + "✅ You're all set with API authentication!", + style="green", + justify="center", + ), + title="🔑 API Key Found", + border_style="bright_green", + ) + console.print(api_key_panel) + console.print() + return False + + # Prompt for authentication method + auth_choices = ["🔐 Login in with Codeflash", "🔑 Use Codeflash API key"] + + questions = [ + inquirer.List( + "auth_method", + message="How would you like to authenticate?", + choices=auth_choices, + default=auth_choices[0], + carousel=True, + ) + ] + + answers = inquirer.prompt(questions, theme=CodeflashTheme()) + if not answers: + apologize_and_exit() + + method = answers["auth_method"] + + if method == auth_choices[1]: + enter_api_key_and_save_to_rc() + ph("cli-new-api-key-entered") + return True + + # Perform OAuth sign-in + api_key = perform_oauth_signin() + + if not api_key: + apologize_and_exit() + + # Save API key + shell_rc_path = get_shell_rc_path() + if not shell_rc_path.exists() and os.name == "nt": + shell_rc_path.touch() + click.echo(f"✅ Created {shell_rc_path}") + + result = save_api_key_to_rc(api_key) + if is_successful(result): + click.echo(result.unwrap()) + click.echo("✅ Signed in successfully and API key saved!") + else: + click.echo(result.failure()) + click.pause() + + os.environ["CODEFLASH_API_KEY"] = api_key + ph("cli-oauth-signin-completed") + return True + + +def enter_api_key_and_save_to_rc() -> None: + browser_launched = False + api_key = "" + while api_key == "": + api_key = click.prompt( + f"Enter your Codeflash API key{' [or press Enter to open your API key page]' if not browser_launched else ''}", + hide_input=False, + default="", + type=CFAPIKeyType(), + show_default=False, + ).strip() + if api_key: + break + if not browser_launched: + click.echo( + f"Opening your Codeflash API key page. Grab a key from there!{LF}" + "You can also open this link manually: https://app.codeflash.ai/app/apikeys" + ) + click.launch("https://app.codeflash.ai/app/apikeys") + browser_launched = True # This does not work on remote consoles + shell_rc_path = get_shell_rc_path() + if not shell_rc_path.exists() and os.name == "nt": + # On Windows, create the appropriate file (PowerShell .ps1 or CMD .bat) in the user's home directory + shell_rc_path.parent.mkdir(parents=True, exist_ok=True) + shell_rc_path.touch() + click.echo(f"✅ Created {shell_rc_path}") + get_user_id(api_key=api_key) # Used to verify whether the API key is valid. + result = save_api_key_to_rc(api_key) + if is_successful(result): + click.echo(result.unwrap()) + else: + click.echo(result.failure()) + click.pause() + + os.environ["CODEFLASH_API_KEY"] = api_key + + +def install_github_app(git_remote: str) -> None: + try: + git_repo = git.Repo(search_parent_directories=True) + 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): + click.echo( + f"🐙 Looks like you've already installed the Codeflash GitHub app on this repository ({owner}/{repo})! Continuing…" + ) + + else: + try: + click.prompt( + f"Finally, you'll need to install the Codeflash GitHub app by choosing the repository you want to install Codeflash on.{LF}" + f"I will attempt to open the github app page - https://github.com/apps/codeflash-ai/installations/select_target {LF}" + f"Please, press ENTER to open the app installation page{LF}", + default="", + type=click.STRING, + prompt_suffix=">>> ", + show_default=False, + ) + click.launch("https://github.com/apps/codeflash-ai/installations/select_target") + click.prompt( + f"Please, press ENTER once you've finished installing the github app from https://github.com/apps/codeflash-ai/installations/select_target{LF}", + default="", + type=click.STRING, + prompt_suffix=">>> ", + show_default=False, + ) + + count = 2 + while not is_github_app_installed_on_repo(owner, repo, suppress_errors=True): + if count == 0: + click.echo( + f"❌ It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo}.{LF}" + f"You won't be able to create PRs with Codeflash until you install the app.{LF}" + f"In the meantime you can make local only optimizations by using the '--no-pr' flag with codeflash.{LF}" + ) + break + click.prompt( + f"❌ It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo}.{LF}" + f"Please install it from https://github.com/apps/codeflash-ai/installations/select_target {LF}" + f"Please, press ENTER to continue once you've finished installing the github app…{LF}", + default="", + type=click.STRING, + prompt_suffix=">>> ", + show_default=False, + ) + count -= 1 + except (KeyboardInterrupt, EOFError, click.exceptions.Abort): + # leave empty line for the next prompt to be properly rendered + click.echo() diff --git a/codeflash/cli_cmds/init_config.py b/codeflash/cli_cmds/init_config.py new file mode 100644 index 000000000..f3f3fc244 --- /dev/null +++ b/codeflash/cli_cmds/init_config.py @@ -0,0 +1,286 @@ +from __future__ import annotations + +import os +from enum import Enum +from functools import lru_cache +from pathlib import Path +from typing import Any, Optional, Union + +import click +import inquirer +import inquirer.themes +import tomlkit +from pydantic.dataclasses import dataclass + +from codeflash.cli_cmds.cli_common import apologize_and_exit +from codeflash.cli_cmds.console import console +from codeflash.code_utils.compat import LF +from codeflash.code_utils.config_parser import parse_config_file +from codeflash.code_utils.env_utils import check_formatter_installed +from codeflash.lsp.helpers import is_LSP_enabled +from codeflash.telemetry.posthog_cf import ph + + +@dataclass(frozen=True) +class CLISetupInfo: + """Setup info for Python projects.""" + + module_root: str + tests_root: str + benchmarks_root: Union[str, None] + ignore_paths: list[str] + formatter: Union[str, list[str]] + git_remote: str + enable_telemetry: bool + + +@dataclass(frozen=True) +class VsCodeSetupInfo: + """Setup info for VSCode extension initialization.""" + + module_root: str + tests_root: str + formatter: Union[str, list[str]] + + +# Custom theme for better UX +class CodeflashTheme(inquirer.themes.Default): + def __init__(self) -> None: + super().__init__() + self.Question.mark_color = inquirer.themes.term.yellow + self.Question.brackets_color = inquirer.themes.term.bright_blue + self.Question.default_color = inquirer.themes.term.bright_cyan + self.List.selection_color = inquirer.themes.term.bright_blue + self.Checkbox.selection_color = inquirer.themes.term.bright_blue + self.Checkbox.selected_icon = "✅" + self.Checkbox.unselected_icon = "⬜" + + +# common sections between normal mode and lsp mode +class CommonSections(Enum): + module_root = "module_root" + tests_root = "tests_root" + formatter_cmds = "formatter_cmds" + + def get_toml_key(self) -> str: + return self.value.replace("_", "-") + + +@lru_cache(maxsize=1) +def get_valid_subdirs(current_dir: Optional[Path] = None) -> list[str]: + ignore_subdirs = [ + "venv", + "node_modules", + "dist", + "build", + "build_temp", + "build_scripts", + "env", + "logs", + "tmp", + "__pycache__", + ] + path_str = str(current_dir) if current_dir else "." + return [ + d + for d in next(os.walk(path_str))[1] + if not d.startswith(".") and not d.startswith("__") and d not in ignore_subdirs + ] + + +def get_suggestions(section: str) -> tuple[list[str], Optional[str]]: + valid_subdirs = get_valid_subdirs() + if section == CommonSections.module_root: + return [d for d in valid_subdirs if d != "tests"], None + if section == CommonSections.tests_root: + default = "tests" if "tests" in valid_subdirs else None + return valid_subdirs, default + if section == CommonSections.formatter_cmds: + return ["disabled", "ruff", "black"], "disabled" + msg = f"Unknown section: {section}" + raise ValueError(msg) + + +def config_found(pyproject_toml_path: Union[str, Path]) -> tuple[bool, str]: + pyproject_toml_path = Path(pyproject_toml_path) + + if not pyproject_toml_path.exists(): + return False, f"Configuration file not found: {pyproject_toml_path}" + + if not pyproject_toml_path.is_file(): + return False, f"Configuration file is not a file: {pyproject_toml_path}" + + if pyproject_toml_path.suffix != ".toml": + return False, f"Configuration file is not a .toml file: {pyproject_toml_path}" + + return True, "" + + +def is_valid_pyproject_toml(pyproject_toml_path: Union[str, Path]) -> tuple[bool, dict[str, Any] | None, str]: + pyproject_toml_path = Path(pyproject_toml_path) + try: + config, _ = parse_config_file(pyproject_toml_path) + except Exception as e: + return False, None, f"Failed to parse configuration: {e}" + + module_root = config.get("module_root") + if not module_root: + return False, config, "Missing required field: 'module_root'" + + if not Path(module_root).is_dir(): + return False, config, f"Invalid 'module_root': directory does not exist at {module_root}" + + tests_root = config.get("tests_root") + if not tests_root: + return False, config, "Missing required field: 'tests_root'" + + if not Path(tests_root).is_dir(): + return False, config, f"Invalid 'tests_root': directory does not exist at {tests_root}" + + return True, 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. + + If it does, ask the user if they want to re-configure it. + """ + from rich.prompt import Confirm + + pyproject_toml_path = Path.cwd() / "pyproject.toml" + + found, _ = config_found(pyproject_toml_path) + if not found: + return True, None + + valid, config, _message = is_valid_pyproject_toml(pyproject_toml_path) + if not valid: + # needs to be re-configured + return True, None + + return Confirm.ask( + "✅ A valid Codeflash config already exists in this project. Do you want to re-configure it?", + default=False, + show_default=True, + ), config + + +def get_formatter_cmds(formatter: str) -> list[str]: + if formatter == "black": + return ["black $file"] + if formatter == "ruff": + return ["ruff check --exit-zero --fix $file", "ruff format $file"] + if formatter == "other": + click.echo( + "🔧 In pyproject.toml, please replace 'your-formatter' with the command you use to format your code." + ) + return ["your-formatter $file"] + if formatter in {"don't use a formatter", "disabled"}: + return ["disabled"] + if " && " in formatter: + return formatter.split(" && ") + return [formatter] + + +# Create or update the pyproject.toml file with the Codeflash dependency & configuration +def configure_pyproject_toml( + setup_info: Union[VsCodeSetupInfo, CLISetupInfo], config_file: Optional[Path] = None +) -> bool: + for_vscode = isinstance(setup_info, VsCodeSetupInfo) + toml_path = config_file or Path.cwd() / "pyproject.toml" + try: + with toml_path.open(encoding="utf8") as pyproject_file: + pyproject_data = tomlkit.parse(pyproject_file.read()) + except FileNotFoundError: + click.echo( + f"I couldn't find a pyproject.toml in the current directory.{LF}" + f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file." + ) + return False + + codeflash_section = tomlkit.table() + codeflash_section.add(tomlkit.comment("All paths are relative to this pyproject.toml's directory.")) + + if for_vscode: + for section in CommonSections: + if hasattr(setup_info, section.value): + codeflash_section[section.get_toml_key()] = getattr(setup_info, section.value) + else: + codeflash_section["module-root"] = setup_info.module_root + codeflash_section["tests-root"] = setup_info.tests_root + codeflash_section["ignore-paths"] = setup_info.ignore_paths + if not setup_info.enable_telemetry: + codeflash_section["disable-telemetry"] = not setup_info.enable_telemetry + if setup_info.git_remote not in ["", "origin"]: + codeflash_section["git-remote"] = setup_info.git_remote + + formatter = setup_info.formatter + + formatter_cmds = formatter if isinstance(formatter, list) else get_formatter_cmds(formatter) + + check_formatter_installed(formatter_cmds, exit_on_failure=False) + codeflash_section["formatter-cmds"] = formatter_cmds + # Add the 'codeflash' section, ensuring 'tool' section exists + tool_section = pyproject_data.get("tool", tomlkit.table()) + + if for_vscode: + # merge the existing codeflash section, instead of overwriting it + existing_codeflash = tool_section.get("codeflash", tomlkit.table()) + + for key, value in codeflash_section.items(): + existing_codeflash[key] = value + tool_section["codeflash"] = existing_codeflash + else: + tool_section["codeflash"] = codeflash_section + + pyproject_data["tool"] = tool_section + + with toml_path.open("w", encoding="utf8") as pyproject_file: + pyproject_file.write(tomlkit.dumps(pyproject_data)) + click.echo(f"Added Codeflash configuration to {toml_path}") + click.echo() + return True + + +def create_empty_pyproject_toml(pyproject_toml_path: Path) -> None: + ph("cli-create-pyproject-toml") + lsp_mode = is_LSP_enabled() + # Define a minimal pyproject.toml content + new_pyproject_toml = tomlkit.document() + new_pyproject_toml["tool"] = {"codeflash": {}} + try: + pyproject_toml_path.write_text(tomlkit.dumps(new_pyproject_toml), encoding="utf8") + + # Check if the pyproject.toml file was created + if pyproject_toml_path.exists() and not lsp_mode: + from rich.panel import Panel + from rich.text import Text + + success_panel = Panel( + Text( + f"✅ Created a pyproject.toml file at {pyproject_toml_path}\n\n" + "Your project is now ready for Codeflash configuration!", + style="green", + justify="center", + ), + title="🎉 Success!", + border_style="bright_green", + ) + console.print(success_panel) + console.print("\n📍 Press any key to continue...") + console.input() + ph("cli-created-pyproject-toml") + except OSError: + click.echo("❌ Failed to create pyproject.toml. Please check your disk permissions and available space.") + apologize_and_exit() + + +def ask_for_telemetry() -> bool: + """Prompt the user to enable or disable telemetry.""" + from rich.prompt import Confirm + + return Confirm.ask( + "⚡️ Help us improve Codeflash by sharing anonymous usage data (e.g. errors encountered)?", + default=True, + show_default=True, + ) diff --git a/codeflash/cli_cmds/init_javascript.py b/codeflash/cli_cmds/init_javascript.py index 5e9f11a47..ed3e50827 100644 --- a/codeflash/cli_cmds/init_javascript.py +++ b/codeflash/cli_cmds/init_javascript.py @@ -69,7 +69,7 @@ class JSSetupInfo: # Import theme from cmd_init to avoid duplication def _get_theme(): """Get the CodeflashTheme - imported lazily to avoid circular imports.""" - from codeflash.cli_cmds.cmd_init import CodeflashTheme + from codeflash.cli_cmds.init_config import CodeflashTheme return CodeflashTheme() @@ -210,7 +210,8 @@ def get_package_install_command(project_root: Path, package: str, dev: bool = Tr def init_js_project(language: ProjectLanguage, *, skip_confirm: bool = False, skip_api_key: bool = False) -> None: """Initialize Codeflash for a JavaScript/TypeScript project.""" - from codeflash.cli_cmds.cmd_init import install_github_actions, install_github_app, prompt_api_key + from codeflash.cli_cmds.github_workflow import install_github_actions + from codeflash.cli_cmds.init_auth import install_github_app, prompt_api_key lang_name = "TypeScript" if language == ProjectLanguage.TYPESCRIPT else "JavaScript" @@ -325,7 +326,7 @@ def collect_js_setup_info(language: ProjectLanguage, *, skip_confirm: bool = Fal Uses auto-detection for most settings and only asks for overrides if needed. When skip_confirm is True, uses all auto-detected defaults without prompting. """ - from codeflash.cli_cmds.cmd_init import ask_for_telemetry, get_valid_subdirs + from codeflash.cli_cmds.init_config import ask_for_telemetry, get_valid_subdirs from codeflash.code_utils.config_js import ( detect_formatter, detect_module_root, diff --git a/codeflash/code_utils/oauth_handler.py b/codeflash/cli_cmds/oauth_handler.py similarity index 100% rename from codeflash/code_utils/oauth_handler.py rename to codeflash/cli_cmds/oauth_handler.py diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py index 32c0ae10c..a9c3558d7 100644 --- a/codeflash/languages/base.py +++ b/codeflash/languages/base.py @@ -678,7 +678,7 @@ class LanguageSupport(Protocol): @property def function_optimizer_class(self) -> type: """Return the FunctionOptimizer subclass for this language.""" - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.function_optimizer import FunctionOptimizer return FunctionOptimizer diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/languages/function_optimizer.py similarity index 100% rename from codeflash/optimization/function_optimizer.py rename to codeflash/languages/function_optimizer.py diff --git a/codeflash/languages/javascript/function_optimizer.py b/codeflash/languages/javascript/function_optimizer.py index 59cf502ce..3bc53965a 100644 --- a/codeflash/languages/javascript/function_optimizer.py +++ b/codeflash/languages/javascript/function_optimizer.py @@ -15,6 +15,7 @@ from codeflash.code_utils.config_consts import ( TOTAL_LOOPING_TIME_EFFECTIVE, ) from codeflash.either import Failure, Success +from codeflash.languages.function_optimizer import FunctionOptimizer from codeflash.models.models import ( CodeOptimizationContext, CodeString, @@ -23,7 +24,6 @@ from codeflash.models.models import ( TestingMode, TestResults, ) -from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.verification.equivalence import compare_test_results if TYPE_CHECKING: diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index d73b85338..830b6c292 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -37,7 +37,6 @@ from codeflash.models.models import ( CodeStringsMarkdown, FunctionSource, ) -from codeflash.optimization.function_context import belongs_to_function_qualified if TYPE_CHECKING: from jedi.api.classes import Name @@ -1470,3 +1469,41 @@ def prune_cst( include_init_dunder=include_init_dunder, ), ) + + +def belongs_to_method(name: Name, class_name: str, method_name: str) -> bool: + """Check if the given name belongs to the specified method.""" + return belongs_to_function(name, method_name) and belongs_to_class(name, class_name) + + +def belongs_to_function(name: Name, function_name: str) -> bool: + """Check if the given jedi Name is a direct child of the specified function.""" + if name.name == function_name: # Handles function definition and recursive function calls + return False + if (name := name.parent()) and name.type == "function": + return name.name == function_name + return False + + +def belongs_to_class(name: Name, class_name: str) -> bool: + """Check if given jedi Name is a direct child of the specified class.""" + while name := name.parent(): + if name.type == "class": + return name.name == class_name + return False + + +def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> bool: + """Check if the given jedi Name is a direct child of the specified function, matched by qualified function name.""" + try: + if ( + name.full_name.startswith(name.module_name) + and get_qualified_name(name.module_name, name.full_name) == qualified_function_name + ): + # Handles function definition and recursive function calls + return False + if (name := name.parent()) and name.type == "function": + return get_qualified_name(name.module_name, name.full_name) == qualified_function_name + return False + except ValueError: + return False diff --git a/codeflash/languages/python/function_optimizer.py b/codeflash/languages/python/function_optimizer.py index 32cfb9187..1677bf8bb 100644 --- a/codeflash/languages/python/function_optimizer.py +++ b/codeflash/languages/python/function_optimizer.py @@ -10,6 +10,7 @@ from codeflash.cli_cmds.console import code_print, console, logger from codeflash.code_utils.code_utils import unified_diff_strings from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE from codeflash.either import Failure, Success +from codeflash.languages.function_optimizer import FunctionOptimizer from codeflash.languages.python.context.unused_definition_remover import ( detect_unused_helper_functions, revert_unused_helper_functions, @@ -22,7 +23,6 @@ from codeflash.languages.python.static_analysis.code_replacer import ( ) from codeflash.languages.python.static_analysis.line_profile_utils import add_decorator_imports, contains_jit_decorator from codeflash.models.models import TestingMode, TestResults -from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.verification.parse_test_output import calculate_function_throughput_from_test_results if TYPE_CHECKING: @@ -87,7 +87,7 @@ class PythonFunctionOptimizer(FunctionOptimizer): return original_conftest_content def instrument_capture(self, file_path_to_helper_classes: dict[Path, set[str]]) -> None: - from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture + from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture instrument_codeflash_capture(self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root) @@ -181,7 +181,7 @@ class PythonFunctionOptimizer(FunctionOptimizer): def parse_line_profile_test_results( self, line_profiler_output_file: Path | None ) -> tuple[TestResults | dict, CoverageData | None]: - from codeflash.verification.parse_line_profile_test_output import parse_line_profile_results + from codeflash.languages.python.parse_line_profile_test_output import parse_line_profile_results return parse_line_profile_results(line_profiler_output_file=line_profiler_output_file) diff --git a/codeflash/verification/instrument_codeflash_capture.py b/codeflash/languages/python/instrument_codeflash_capture.py similarity index 100% rename from codeflash/verification/instrument_codeflash_capture.py rename to codeflash/languages/python/instrument_codeflash_capture.py diff --git a/codeflash/verification/parse_line_profile_test_output.py b/codeflash/languages/python/parse_line_profile_test_output.py similarity index 100% rename from codeflash/verification/parse_line_profile_test_output.py rename to codeflash/languages/python/parse_line_profile_test_output.py diff --git a/codeflash/languages/python/reference_graph.py b/codeflash/languages/python/reference_graph.py index 4f389fd66..46867a74d 100644 --- a/codeflash/languages/python/reference_graph.py +++ b/codeflash/languages/python/reference_graph.py @@ -79,7 +79,7 @@ def _is_valid_definition(definition: Name, caller_qualified_name: str, project_r return False try: - from codeflash.optimization.function_context import belongs_to_function_qualified + from codeflash.languages.python.context.code_context_extractor import belongs_to_function_qualified if belongs_to_function_qualified(definition, caller_qualified_name): return False diff --git a/codeflash/languages/python/support.py b/codeflash/languages/python/support.py index 2652fcf83..c6400913b 100644 --- a/codeflash/languages/python/support.py +++ b/codeflash/languages/python/support.py @@ -1051,8 +1051,8 @@ class PythonSupport: from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE from codeflash.languages.python.static_analysis.coverage_utils import prepare_coverage_files + from codeflash.languages.python.test_runner import execute_test_subprocess from codeflash.models.models import TestType - from codeflash.verification.test_runner import execute_test_subprocess blocklisted_plugins = ["benchmark", "codspeed", "xdist", "sugar"] @@ -1156,7 +1156,7 @@ class PythonSupport: from codeflash.code_utils.code_utils import get_run_tmp_file from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE - from codeflash.verification.test_runner import execute_test_subprocess + from codeflash.languages.python.test_runner import execute_test_subprocess blocklisted_plugins = ["codspeed", "cov", "benchmark", "profiling", "xdist", "sugar"] @@ -1200,7 +1200,7 @@ class PythonSupport: from codeflash.code_utils.code_utils import get_run_tmp_file from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME_EFFECTIVE - from codeflash.verification.test_runner import execute_test_subprocess + from codeflash.languages.python.test_runner import execute_test_subprocess blocklisted_plugins = ["codspeed", "cov", "benchmark", "profiling", "xdist", "sugar"] diff --git a/codeflash/verification/test_runner.py b/codeflash/languages/python/test_runner.py similarity index 100% rename from codeflash/verification/test_runner.py rename to codeflash/languages/python/test_runner.py diff --git a/codeflash/lsp/beta.py b/codeflash/lsp/beta.py index 68d332b21..31349b841 100644 --- a/codeflash/lsp/beta.py +++ b/codeflash/lsp/beta.py @@ -11,13 +11,13 @@ from typing import TYPE_CHECKING, Optional, Union from codeflash.api.cfapi import get_codeflash_api_key, get_user_id from codeflash.cli_cmds.cli import process_pyproject_config -from codeflash.cli_cmds.cmd_init import ( +from codeflash.cli_cmds.cmd_init import create_find_common_tags_file +from codeflash.cli_cmds.init_config import ( CommonSections, VsCodeSetupInfo, config_found, configure_pyproject_toml, create_empty_pyproject_toml, - create_find_common_tags_file, get_formatter_cmds, get_suggestions, get_valid_subdirs, diff --git a/codeflash/main.py b/codeflash/main.py index 32ae9c66c..7b2a557db 100644 --- a/codeflash/main.py +++ b/codeflash/main.py @@ -18,14 +18,11 @@ if "--subagent" in sys.argv: warnings.filterwarnings("ignore") from codeflash.cli_cmds.cli import parse_args, process_pyproject_config -from codeflash.cli_cmds.cmd_init import CODEFLASH_LOGO, ask_run_end_to_end_test from codeflash.cli_cmds.console import paneled_text from codeflash.code_utils import env_utils from codeflash.code_utils.checkpoint import ask_should_use_checkpoint_get_functions from codeflash.code_utils.config_parser import parse_config_file from codeflash.code_utils.version_check import check_for_newer_minor_version -from codeflash.telemetry import posthog_cf -from codeflash.telemetry.sentry import init_sentry if TYPE_CHECKING: from argparse import Namespace @@ -33,6 +30,9 @@ if TYPE_CHECKING: def main() -> None: """Entry point for the codeflash command-line interface.""" + from codeflash.telemetry import posthog_cf + from codeflash.telemetry.sentry import init_sentry + args = parse_args() print_codeflash_banner() @@ -46,11 +46,30 @@ def main() -> None: disable_telemetry = pyproject_config.get("disable_telemetry", False) init_sentry(enabled=not disable_telemetry, exclude_errors=True) posthog_cf.initialize_posthog(enabled=not disable_telemetry) - args.func() + + if args.command == "init": + from codeflash.cli_cmds.cmd_init import init_codeflash + + init_codeflash() + elif args.command == "init-actions": + from codeflash.cli_cmds.github_workflow import install_github_actions + + install_github_actions() + elif args.command == "vscode-install": + from codeflash.cli_cmds.extension import install_vscode_extension + + install_vscode_extension() + elif args.command == "optimize": + from codeflash.tracer import main as tracer_main + + tracer_main(args) elif args.verify_setup: args = process_pyproject_config(args) init_sentry(enabled=not args.disable_telemetry, exclude_errors=True) posthog_cf.initialize_posthog(enabled=not args.disable_telemetry) + + from codeflash.cli_cmds.cmd_init import ask_run_end_to_end_test + ask_run_end_to_end_test(args) else: # Check for first-run experience (no config exists) @@ -117,6 +136,8 @@ def _handle_config_loading(args: Namespace) -> Namespace | None: def print_codeflash_banner() -> None: + from codeflash.cli_cmds.console_constants import CODEFLASH_LOGO + paneled_text( CODEFLASH_LOGO, panel_args={"title": "https://codeflash.ai", "expand": False}, text_args={"style": "bold gold3"} ) diff --git a/codeflash/optimization/function_context.py b/codeflash/optimization/function_context.py deleted file mode 100644 index eb710945f..000000000 --- a/codeflash/optimization/function_context.py +++ /dev/null @@ -1,46 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -from codeflash.code_utils.code_utils import get_qualified_name - -if TYPE_CHECKING: - from jedi.api.classes import Name - - -def belongs_to_method(name: Name, class_name: str, method_name: str) -> bool: - """Check if the given name belongs to the specified method.""" - return belongs_to_function(name, method_name) and belongs_to_class(name, class_name) - - -def belongs_to_function(name: Name, function_name: str) -> bool: - """Check if the given jedi Name is a direct child of the specified function.""" - if name.name == function_name: # Handles function definition and recursive function calls - return False - if (name := name.parent()) and name.type == "function": - return name.name == function_name - return False - - -def belongs_to_class(name: Name, class_name: str) -> bool: - """Check if given jedi Name is a direct child of the specified class.""" - while name := name.parent(): - if name.type == "class": - return name.name == class_name - return False - - -def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> bool: - """Check if the given jedi Name is a direct child of the specified function, matched by qualified function name.""" - try: - if ( - name.full_name.startswith(name.module_name) - and get_qualified_name(name.module_name, name.full_name) == qualified_function_name - ): - # Handles function definition and recursive function calls - return False - if (name := name.parent()) and name.type == "function": - return get_qualified_name(name.module_name, name.full_name) == qualified_function_name - return False - except ValueError: - return False diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index cd8e41bbe..21fe83ff2 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -42,8 +42,8 @@ if TYPE_CHECKING: from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.languages.base import DependencyResolver + from codeflash.languages.function_optimizer import FunctionOptimizer from codeflash.models.models import BenchmarkKey, FunctionCalledInTest, ValidCode - from codeflash.optimization.function_optimizer import FunctionOptimizer class Optimizer: diff --git a/codeflash/tracer.py b/codeflash/tracer.py index 3f1bde3d1..3826eca15 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -199,8 +199,8 @@ def main(args: Namespace | None = None) -> ArgumentParser: result_pickle_file_path.unlink(missing_ok=True) if not parsed_args.trace_only and replay_test_paths: from codeflash.cli_cmds.cli import parse_args, process_pyproject_config - from codeflash.cli_cmds.cmd_init import CODEFLASH_LOGO from codeflash.cli_cmds.console import paneled_text + from codeflash.cli_cmds.console_constants import CODEFLASH_LOGO from codeflash.languages import set_current_language from codeflash.languages.base import Language from codeflash.telemetry import posthog_cf diff --git a/tests/benchmarks/test_benchmark_code_extract_code_context.py b/tests/benchmarks/test_benchmark_code_extract_code_context.py index 77c435720..4fe06b14d 100644 --- a/tests/benchmarks/test_benchmark_code_extract_code_context.py +++ b/tests/benchmarks/test_benchmark_code_extract_code_context.py @@ -22,7 +22,7 @@ def test_benchmark_extract(benchmark) -> None: ) function_to_optimize = FunctionToOptimize( function_name="replace_function_and_helpers_with_optimized_code", - file_path=file_path / "optimization" / "function_optimizer.py", + file_path=file_path / "languages" / "function_optimizer.py", parents=[FunctionParent(name="FunctionOptimizer", type="ClassDef")], starting_line=None, ending_line=None, diff --git a/tests/test_async_run_and_parse_tests.py b/tests/test_async_run_and_parse_tests.py index 1777a1c73..1034b5e51 100644 --- a/tests/test_async_run_and_parse_tests.py +++ b/tests/test_async_run_and_parse_tests.py @@ -16,7 +16,7 @@ from codeflash.code_utils.instrument_existing_tests import ( from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestType from codeflash.optimization.optimizer import Optimizer -from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture +from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") diff --git a/tests/test_cmd_init.py b/tests/test_cmd_init.py index ee456783b..1e3523681 100644 --- a/tests/test_cmd_init.py +++ b/tests/test_cmd_init.py @@ -4,7 +4,7 @@ from pathlib import Path import pytest -from codeflash.cli_cmds.cmd_init import ( +from codeflash.cli_cmds.init_config import ( CLISetupInfo, VsCodeSetupInfo, configure_pyproject_toml, diff --git a/tests/test_codeflash_capture.py b/tests/test_codeflash_capture.py index ddf44d017..21d27fb4c 100644 --- a/tests/test_codeflash_capture.py +++ b/tests/test_codeflash_capture.py @@ -10,8 +10,8 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import FunctionParent, TestFile, TestFiles, TestingMode, TestType, VerificationType from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer from codeflash.verification.equivalence import compare_test_results -from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture -from codeflash.verification.test_runner import execute_test_subprocess +from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture +from codeflash.languages.python.test_runner import execute_test_subprocess from codeflash.verification.verification_utils import TestConfig diff --git a/tests/test_formatter.py b/tests/test_formatter.py index b7eee0f52..a3998a81f 100644 --- a/tests/test_formatter.py +++ b/tests/test_formatter.py @@ -9,7 +9,7 @@ from codeflash.code_utils.config_parser import parse_config_file from codeflash.code_utils.formatter import format_code, format_generated_code, sort_imports from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import CodeString, CodeStringsMarkdown -from codeflash.optimization.function_optimizer import FunctionOptimizer +from codeflash.languages.function_optimizer import FunctionOptimizer from codeflash.verification.verification_utils import TestConfig diff --git a/tests/test_instrument_all_and_run.py b/tests/test_instrument_all_and_run.py index a8ed56f15..40fb8fbcf 100644 --- a/tests/test_instrument_all_and_run.py +++ b/tests/test_instrument_all_and_run.py @@ -12,7 +12,7 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestType from codeflash.optimization.optimizer import Optimizer from codeflash.verification.equivalence import compare_test_results -from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture +from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture # Used by cli instrumentation codeflash_wrap_string = """def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, codeflash_cur, codeflash_con, *args, **kwargs): diff --git a/tests/test_instrument_codeflash_capture.py b/tests/test_instrument_codeflash_capture.py index 8228cc488..543d50855 100644 --- a/tests/test_instrument_codeflash_capture.py +++ b/tests/test_instrument_codeflash_capture.py @@ -3,7 +3,7 @@ from pathlib import Path from codeflash.code_utils.code_utils import get_run_tmp_file from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import FunctionParent -from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture +from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture def test_add_codeflash_capture(): diff --git a/tests/test_instrumentation_run_results_aiservice.py b/tests/test_instrumentation_run_results_aiservice.py index 4879cc93a..69fa82d2e 100644 --- a/tests/test_instrumentation_run_results_aiservice.py +++ b/tests/test_instrumentation_run_results_aiservice.py @@ -14,7 +14,7 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import FunctionParent, TestFile, TestFiles, TestingMode, TestType, VerificationType from codeflash.optimization.optimizer import Optimizer from codeflash.verification.equivalence import compare_test_results -from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture +from codeflash.languages.python.instrument_codeflash_capture import instrument_codeflash_capture # Used by aiservice instrumentation behavior_logging_code = """ diff --git a/tests/test_languages/test_javascript_optimization_flow.py b/tests/test_languages/test_javascript_optimization_flow.py index 22c2ab6bc..844b8d683 100644 --- a/tests/test_languages/test_javascript_optimization_flow.py +++ b/tests/test_languages/test_javascript_optimization_flow.py @@ -305,7 +305,7 @@ describe('fibonacci', () => { """Test FunctionOptimizer can be instantiated for JavaScript.""" skip_if_js_not_supported() from codeflash.discovery.functions_to_optimize import find_all_functions_in_file - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.function_optimizer import FunctionOptimizer src_file = js_project / "utils.js" functions = find_all_functions_in_file(src_file) @@ -340,7 +340,7 @@ describe('fibonacci', () => { """Test FunctionOptimizer can be instantiated for TypeScript.""" skip_if_js_not_supported() from codeflash.discovery.functions_to_optimize import find_all_functions_in_file - from codeflash.optimization.function_optimizer import FunctionOptimizer + from codeflash.languages.function_optimizer import FunctionOptimizer src_file = ts_project / "utils.ts" functions = find_all_functions_in_file(src_file)