codeflash 0.3.4__tar.gz → 0.3.6__tar.gz

{codeflash-0.3.4 → codeflash-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: codeflash
-Version: 0.3.4
+Version: 0.3.6
 Summary: A Code Performance Optimization Library
 Home-page: https://codeflash.ai
 Author: CodeFlash Inc.
@@ -19,10 +19,12 @@ Requires-Dist: inquirer (>=3.0.0)
 Requires-Dist: jedi (>=0.19.0)
 Requires-Dist: junitparser (>=3.1.0)
 Requires-Dist: libcst (>=1.0.1)
+Requires-Dist: parameterized (>=0.9.0)
 Requires-Dist: posthog (>=3.0.0)
 Requires-Dist: pydantic (>=1.10.1)
 Requires-Dist: pytest (>=7.0.0)
 Requires-Dist: pytest-timeout (>=2.1.0)
+Requires-Dist: sentry-sdk (>=1.40.6,<2.0.0)
 Requires-Dist: tiktoken (>=0.3.2)
 Requires-Dist: timeout-decorator (>=0.5.0)
 Requires-Dist: tomlkit (>=0.11.7)

{codeflash-0.3.4 → codeflash-0.3.6}/codeflash/analytics/posthog.py

@@ -11,18 +11,33 @@ _posthog = Posthog(
 )
 
 
+_ANALYTICS_ENABLED = True
+
+
+def enable_analytics(enabled: bool) -> None:
+    """
+    Enable or disable analytics.
+    :param enabled: Whether to enable analytics.
+    """
+    if enabled:
+        ph("cli-analytics-enabled")
+    else:
+        ph("cli-analytics-disabled")
+    global _ANALYTICS_ENABLED
+    _ANALYTICS_ENABLED = enabled
+
+
 def ph(event: str, properties: Dict[str, Any] = None) -> None:
     """
     Log an event to PostHog.
     :param event: The name of the event.
     :param properties: A dictionary of properties to attach to the event.
     """
+    if not _ANALYTICS_ENABLED:
+        return
 
-    if properties is None:
-        properties = {}
-
-    properties["cli_version"] = __version__
-    properties["cli_version_tuple"] = __version_tuple__
+    properties = properties or {}
+    properties.update({"cli_version": __version__, "cli_version_tuple": __version_tuple__})
 
     user_id = get_user_id()
 

codeflash-0.3.6/codeflash/analytics/sentry.py

@@ -0,0 +1,14 @@
+import sentry_sdk
+
+
+def init_sentry():
+    sentry_sdk.init(
+        dsn="https://4b9a1902f9361b48c04376df6483bc96@o4506833230561280.ingest.sentry.io/4506833262477312",
+        # Set traces_sample_rate to 1.0 to capture 100%
+        # of transactions for performance monitoring.
+        traces_sample_rate=1.0,
+        # Set profiles_sample_rate to 1.0 to profile 100%
+        # of sampled transactions.
+        # We recommend adjusting this value in production.
+        profiles_sample_rate=1.0,
+    )

{codeflash-0.3.4 → codeflash-0.3.6}/codeflash/api/cfapi.py

@@ -1,15 +1,14 @@
 import json
 import logging
 import os
-from functools import lru_cache
-from typing import Optional, Dict, Any
-
 import requests
+from functools import lru_cache
 from pydantic.json import pydantic_encoder
 from requests import Response
+from typing import Optional, Dict, Any
 
 from codeflash.code_utils.env_utils import get_codeflash_api_key
-from codeflash.github.PrComment import PrComment
+from codeflash.github.PrComment import PrComment, FileDiffContent
 
 if os.environ.get("CFAPI_SERVER", default="prod").lower() == "local":
     CFAPI_BASE_URL = "http://localhost:3001"
@@ -59,8 +58,9 @@ def suggest_changes(
     owner: str,
     repo: str,
     pr_number: int,
-    file_changes: dict[str, dict[str, str]],
+    file_changes: dict[str, FileDiffContent],
     pr_comment: PrComment,
+    existing_tests: str,
     generated_tests: str,
 ) -> Response:
     """
@@ -81,6 +81,7 @@ def suggest_changes(
         "pullNumber": pr_number,
         "diffContents": file_changes,
         "prCommentFields": pr_comment.to_json(),
+        "existingTests": existing_tests,
         "generatedTests": generated_tests,
     }
     response = make_cfapi_request(endpoint="/suggest-pr-changes", method="POST", payload=payload)
@@ -91,8 +92,9 @@ def create_pr(
     owner: str,
     repo: str,
     base_branch: str,
-    file_changes: dict[str, dict[str, str]],
+    file_changes: dict[str, FileDiffContent],
     pr_comment: PrComment,
+    existing_tests: str,
     generated_tests: str,
 ) -> Response:
     """
@@ -111,6 +113,7 @@ def create_pr(
         "baseBranch": base_branch,
         "diffContents": file_changes,
         "prCommentFields": pr_comment.to_json(),
+        "existingTests": existing_tests,
         "generatedTests": generated_tests,
     }
     response = make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload)
@@ -124,8 +127,7 @@ def check_github_app_installed_on_repo(owner: str, repo: str) -> Response:
     :param repo: The name of the repository.
     :return: The response object.
     """
-    response = make_cfapi_request(
+    return make_cfapi_request(
         endpoint=f"/is-github-app-installed?repo={repo}&owner={owner}",
         method="GET",
     )
-    return response

{codeflash-0.3.4 → codeflash-0.3.6}/codeflash/cli_cmds/cli.py

@@ -11,7 +11,6 @@ from codeflash.cli_cmds.logging_config import LOGGING_FORMAT
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.code_utils.git_utils import (
-    git_root_dir,
     get_repo_owner_and_name,
     get_github_secrets_page_url,
 )
@@ -34,7 +33,11 @@ def process_cmd_args(args: Namespace) -> Namespace:
             raise ValueError(f"File {args.file} does not exist")
         args.file = os.path.realpath(args.file)
 
-    pyproject_config = parse_config_file(args.config_file)
+    try:
+        pyproject_config = parse_config_file(args.config_file)
+    except ValueError as e:
+        logging.error(e.args[0])
+        exit(1)
     supported_keys = [
         "module_root",
         "tests_root",
@@ -42,6 +45,8 @@ def process_cmd_args(args: Namespace) -> Namespace:
         "ignore_paths",
         "minimum_performance_gain",
         "pytest_cmd",
+        "formatter_cmd",
+        "enable_analytics",
     ]
     for key in supported_keys:
         if key in pyproject_config:
@@ -82,15 +87,14 @@ def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
     if hasattr(args, "all"):
         # Ensure that the user can actually open PRs on the repo.
         try:
-            repo = git.Repo(search_parent_directories=True)
-            git_root_dir(repo)
+            git_repo = git.Repo(search_parent_directories=True)
         except git.exc.InvalidGitRepositoryError:
             logging.error(
-                "Could not find a git repository in the current directory. "
-                "We need a git repository to run --all and open PRs for optimizations. Exiting..."
+                "I couldn't find a git repository in the current directory. "
+                "I need a git repository to run --all and open PRs for optimizations. Exiting..."
             )
             exit(1)
-        owner, repo = get_repo_owner_and_name(repo)
+        owner, repo = get_repo_owner_and_name(git_repo)
         try:
             response = check_github_app_installed_on_repo(owner, repo)
             if response.ok and response.text == "true":
@@ -100,7 +104,8 @@ def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
                 raise Exception
         except Exception as e:
             logging.error(
-                f"Could not find the CodeFlash GitHub App installed on the repository {owner}/{repo}. "
+                f"Could not find the CodeFlash GitHub App installed on the repository {owner}/{repo} or the GitHub"
+                f" account linked to your CODEFLASH_API_KEY does not have access to the repository {owner}/{repo}. "
                 "Please install the CodeFlash GitHub App on your repository to use --all."
                 " Instructions at https://app.codeflash.ai \n"
                 "Exiting..."

{codeflash-0.3.4 → codeflash-0.3.6}/codeflash/cli_cmds/cmd_init.py

@@ -8,11 +8,17 @@ from typing import Optional
 
 import click
 import inquirer
+import inquirer.themes
 import tomlkit
 from git import Repo
 
 from codeflash.analytics.posthog import ph
-from codeflash.code_utils.env_utils import get_codeflash_api_key
+from codeflash.code_utils.env_utils import (
+    get_codeflash_api_key,
+    read_api_key_from_shell_config,
+    SHELL_RC_EXPORT_PATTERN,
+)
+from codeflash.code_utils.env_utils import get_shell_rc_path
 from codeflash.code_utils.git_utils import get_github_secrets_page_url
 from codeflash.version import __version__ as version
 
@@ -58,6 +64,8 @@ def init_codeflash():
         click.echo(
             "🐚 Don't forget to restart your shell to load the CODEFLASH_API_KEY environment variable!"
         )
+        click.echo("Or run the following command to reload:")
+        click.echo(f"  source {get_shell_rc_path()}")
 
     ph("cli-installation-successful", {"did_add_new_key": did_add_new_key})
 
@@ -72,7 +80,7 @@ def ask_run_end_to_end_test(setup_info):
             )
         ]
     )
-    run_tests = run_tests_answer["run_tests"]
+    run_tests = run_tests_answer.get("run_tests", False)
     if run_tests:
         create_bubble_sort_file(setup_info)
         run_end_to_end_test(setup_info)
@@ -120,9 +128,11 @@ def collect_setup_info(setup_info: dict[str, str]):
                 message="Which Python module do you want me to optimize going forward?\n"
                 + "(This is usually the top-most directory where all your Python source code is located)",
                 choices=module_subdir_options,
-                default=project_name
-                if project_name in module_subdir_options
-                else module_subdir_options[0],
+                default=(
+                    project_name
+                    if project_name in module_subdir_options
+                    else module_subdir_options[0]
+                ),
             )
         ]
     )
@@ -134,6 +144,8 @@ def collect_setup_info(setup_info: dict[str, str]):
     default_tests_subdir = "tests"
     create_for_me_option = "okay, create a tests/ directory for me!"
     test_subdir_options = valid_subdirs if len(valid_subdirs) > 0 else [create_for_me_option]
+    custom_dir_option = "enter a custom directory..."
+    test_subdir_options.append(custom_dir_option)
     tests_root_answer = inquirer.prompt(
         [
             inquirer.List(
@@ -141,9 +153,11 @@ def collect_setup_info(setup_info: dict[str, str]):
                 message="Where are your tests located? "
                 "(If you don't have any tests yet, I can create an empty tests/ directory for you)",
                 choices=test_subdir_options,
-                default=default_tests_subdir
-                if default_tests_subdir in test_subdir_options
-                else test_subdir_options[0],
+                default=(
+                    default_tests_subdir
+                    if default_tests_subdir in test_subdir_options
+                    else test_subdir_options[0]
+                ),
             )
         ]
     )
@@ -151,7 +165,20 @@ def collect_setup_info(setup_info: dict[str, str]):
     if tests_root == create_for_me_option:
         tests_root = os.path.join(curdir, default_tests_subdir)
         os.mkdir(tests_root)
-        click.echo(f"✅ Created directory {tests_root}/")
+        click.echo(f"✅ Created directory {tests_root}/\n")
+    elif tests_root == custom_dir_option:
+        custom_tests_root_answer = inquirer.prompt(
+            [
+                inquirer.Path(
+                    "custom_tests_root",  # Removed the colon and space from the message
+                    message=f"Enter the path to your tests directory inside {os.path.abspath(module_root) + os.sep} ",
+                    path_type=inquirer.Path.DIRECTORY,
+                    exists=True,
+                    normalize_to_absolute_path=True,
+                ),
+            ]
+        )
+        tests_root = custom_tests_root_answer["custom_tests_root"]
     setup_info["tests_root"] = os.path.relpath(tests_root, curdir)
     ph("cli-tests-root-provided")
 
@@ -179,22 +206,38 @@ def collect_setup_info(setup_info: dict[str, str]):
     ignore_paths = []
     setup_info["ignore_paths"] = ignore_paths
 
+    # Ask the user if they agree to enable PostHog analytics logging
+    # enable_analytics_question = [
+    #     inquirer.List(
+    #         "enable_analytics",
+    #         message="⚡️ Is it OK to collect usage analytics to help improve CodeFlash? (recommended)",
+    #         choices=[
+    #             ("Sure, I'd love to help make CodeFlash better!", True),
+    #             ("No, thanks.", False),
+    #         ],
+    #     )
+    # ]
+    # enable_analytics_answer = inquirer.prompt(enable_analytics_question)
+    # setup_info["enable_analytics"] = enable_analytics_answer["enable_analytics"]
+
+    ph("cli-analytics-choice", {"enable_analytics": setup_info["enable_analytics"]})
+
 
 def detect_test_framework(curdir, tests_root) -> Optional[str]:
     test_framework = None
     pytest_files = ["pytest.ini", "pyproject.toml", "tox.ini", "setup.cfg"]
     pytest_config_patterns = {
-        "pytest.ini": r"\[pytest\]",
-        "pyproject.toml": r"\[tool\.pytest\.ini_options\]",
-        "tox.ini": r"\[pytest\]",
-        "setup.cfg": r"\[tool:pytest\]",
+        "pytest.ini": "[pytest]",
+        "pyproject.toml": "[tool.pytest.ini_options]",
+        "tox.ini": "[pytest]",
+        "setup.cfg": "[tool:pytest]",
     }
     for pytest_file in pytest_files:
         file_path = os.path.join(curdir, pytest_file)
         if os.path.exists(file_path):
-            with open(file_path, "r") as file:
+            with open(file_path, "r", encoding="utf8") as file:
                 contents = file.read()
-                if re.search(pytest_config_patterns[pytest_file], contents):
+                if pytest_config_patterns[pytest_file] in contents:
                     test_framework = "pytest"
                     break
         test_framework = "pytest"
@@ -202,9 +245,12 @@ def detect_test_framework(curdir, tests_root) -> Optional[str]:
         # Check if any python files contain a class that inherits from unittest.TestCase
         for filename in os.listdir(tests_root):
             if filename.endswith(".py"):
-                with open(os.path.join(tests_root, filename), "r") as file:
+                with open(os.path.join(tests_root, filename), "r", encoding="utf8") as file:
                     contents = file.read()
-                    node = ast.parse(contents)
+                    try:
+                        node = ast.parse(contents)
+                    except SyntaxError:
+                        continue
                     if any(
                         isinstance(item, ast.ClassDef)
                         and any(
@@ -230,7 +276,7 @@ def check_for_toml_or_setup_file() -> Optional[str]:
     project_name = None
     if os.path.exists(pyproject_toml_path):
         try:
-            with open(pyproject_toml_path, "r") as f:
+            with open(pyproject_toml_path, "r", encoding="utf8") as f:
                 pyproject_toml_content = f.read()
             project_name = tomlkit.parse(pyproject_toml_content)["tool"]["poetry"]["name"]
             click.echo(f"✅ I found a pyproject.toml for your project {project_name}.")
@@ -239,7 +285,7 @@ def check_for_toml_or_setup_file() -> Optional[str]:
             click.echo(f"✅ I found a pyproject.toml for your project.")
             ph("cli-pyproject-toml-found")
     elif os.path.exists(setup_py_path):
-        with open(setup_py_path, "r") as f:
+        with open(setup_py_path, "r", encoding="utf8") as f:
             setup_py_content = f.read()
         project_name_match = re.search(
             r"setup\s*\([^)]*?name\s*=\s*['\"](.*?)['\"]", setup_py_content, re.DOTALL
@@ -251,38 +297,54 @@ def check_for_toml_or_setup_file() -> Optional[str]:
         else:
             click.echo(f"✅ Found setup.py.")
             ph("cli-setup-py-found")
-        # Create a pyproject.toml file because it doesn't exist
-        create_toml = (
-            click.prompt(
-                f"I need your project to have a pyproject.toml file to store CodeFlash configuration settings.\n"
-                f"Do you want to run `poetry init` to create one?",
-                default="y",
-                type=click.STRING,
-            )
-            .lower()
-            .strip()
-        )
-        if create_toml.startswith("y"):
-            # Check if Poetry is installed, if not, install it using pip
-            poetry_check = subprocess.run(["poetry", "--version"], capture_output=True, text=True)
-            if poetry_check.returncode != 0:
-                click.echo("Poetry is not installed. Installing Poetry...")
-                subprocess.run(["pip", "install", "poetry"], check=True)
-            subprocess.run(["poetry", "init"], cwd=curdir)
-            click.echo(f"✅ Created a pyproject.toml file at {pyproject_toml_path}")
-            ph("cli-created-pyproject-toml")
     else:
         click.echo(
-            f"❌ I couldn't find a pyproject.toml or a setup.py in the current directory ({curdir}).\n"
-            "Please make sure you're running codeflash init from your project's root directory.\n"
-            "See https://app.codeflash.ai/app/getting-started for more details!"
+            f"💡 I couldn't find a pyproject.toml in the current directory ({curdir}).\n"
+            "(make sure you're running `codeflash init` from your project's root directory!)\n"
+            f"I need this file to store my configuration settings."
         )
         ph("cli-no-pyproject-toml-or-setup-py")
-        sys.exit(1)
+
+        # Create a pyproject.toml file because it doesn't exist
+        create_toml = inquirer.confirm(
+            f"Do you want me to create a pyproject.toml file in the current directory?",
+            default=True,
+            show_default=False,
+        )
+        if create_toml:
+            ph("cli-create-pyproject-toml")
+            # Define a minimal pyproject.toml content
+            new_pyproject_toml = tomlkit.document()
+            new_pyproject_toml["tool"] = {"codeflash": {}}
+            try:
+                with open(pyproject_toml_path, "w", encoding="utf8") as pyproject_file:
+                    pyproject_file.write(tomlkit.dumps(new_pyproject_toml))
+
+                # Check if the pyproject.toml file was created
+                if os.path.exists(pyproject_toml_path):
+                    click.echo(f"✅ Created a pyproject.toml file at {pyproject_toml_path}")
+                    click.pause()
+                ph("cli-created-pyproject-toml")
+            except IOError as e:
+                click.echo(
+                    "❌ Failed to create pyproject.toml. Please check your disk permissions and available space."
+                )
+                apologize_and_exit()
+        else:
+            click.echo("⏩️ Skipping pyproject.toml creation.")
+            apologize_and_exit()
     click.echo()
     return project_name
 
 
+def apologize_and_exit():
+    click.echo(
+        "💡 If you're having trouble, see https://app.codeflash.ai/app/getting-started for further help getting started with CodeFlash!"
+    )
+    click.echo("Exiting...")
+    sys.exit(1)
+
+
 # Ask if the user wants CodeFlash to optimize new GitHub PRs
 def prompt_github_action(setup_info: dict[str, str]):
     optimize_prs_answer = inquirer.prompt(
@@ -329,7 +391,7 @@ def prompt_github_action(setup_info: dict[str, str]):
             optimize_yml_content = optimize_yml_content.replace(
                 " {{ python_version }}", python_version_string
             )
-            with open(optimize_yaml_path, "w") as optimize_yml_file:
+            with open(optimize_yaml_path, "w", encoding="utf8") as optimize_yml_file:
                 optimize_yml_file.write(optimize_yml_content)
             click.echo(f"✅ Created {optimize_yaml_path}\n")
             click.prompt(
@@ -376,24 +438,29 @@ def prompt_github_action(setup_info: dict[str, str]):
 def configure_pyproject_toml(setup_info: dict[str, str]):
     toml_path = os.path.join(os.getcwd(), "pyproject.toml")
     try:
-        with open(toml_path, "r") as pyproject_file:
+        with open(toml_path, "r", encoding="utf8") as pyproject_file:
             pyproject_data = tomlkit.parse(pyproject_file.read())
     except FileNotFoundError:
         click.echo(
-            f"Could not find a pyproject.toml in the current directory.\n"
+            f"I couln't find a pyproject.toml in the current directory.\n"
             f"Please create it by running `poetry init`, or run `codeflash init` again from a different project directory."
         )
+        apologize_and_exit()
 
     codeflash_section = tomlkit.table()
     codeflash_section["module-root"] = setup_info["module_root"]
     codeflash_section["tests-root"] = setup_info["tests_root"]
     codeflash_section["test-framework"] = setup_info["test_framework"]
     codeflash_section["ignore-paths"] = setup_info["ignore_paths"]
+    codeflash_section["enable-analytics"] = setup_info["enable_analytics"]
+
+    # Add the 'codeflash' section, ensuring 'tool' section exists
+    tool_section = pyproject_data.get("tool", tomlkit.table())
+    tool_section["codeflash"] = codeflash_section
+    pyproject_data["tool"] = tool_section
 
-    # Add the 'codeflash' section
-    pyproject_data["tool"]["codeflash"] = codeflash_section
     click.echo(f"Writing CodeFlash configuration ...\r", nl=False)
-    with open(toml_path, "w") as pyproject_file:
+    with open(toml_path, "w", encoding="utf8") as pyproject_file:
         pyproject_file.write(tomlkit.dumps(pyproject_data))
     click.echo(f"✅ Added CodeFlash configuration to {toml_path}")
     click.echo()
@@ -407,7 +474,11 @@ class CFAPIKeyType(click.ParamType):
         if value.startswith("cf-") or value == "":
             return value
         else:
-            self.fail(f"{value} does not start with the prefix 'cf-'. Please retry.", param, ctx)
+            self.fail(
+                f"That key [{value}] seems to be invalid. It should start with a 'cf-' prefix. Please try again.",
+                param,
+                ctx,
+            )
 
 
 # Returns True if the user entered a new API key, False if they used an existing one
@@ -418,34 +489,31 @@ def prompt_api_key() -> bool:
         existing_api_key = None
     if existing_api_key:
         display_key = f"{existing_api_key[:3]}****{existing_api_key[-4:]}"
-        use_existing_key = click.prompt(
-            f"I found a CODEFLASH_API_KEY in your environment [{display_key}]!\n"
-            f"Press Enter to use this key, or type any other key to change it",
-            default="",
-            type=CFAPIKeyType(),
+        click.echo(f"🔑 I found a CODEFLASH_API_KEY in your environment [{display_key}]!")
+
+        use_existing_key = inquirer.confirm(
+            message="Do you want to use this key?",
+            default=True,
             show_default=False,
-        ).strip()
-        if use_existing_key == "":
+        )
+        if use_existing_key:
             ph("cli-existing-api-key-used")
             return False
-        else:
-            enter_api_key_and_save_to_rc(existing_api_key=use_existing_key)
-            ph("cli-new-api-key-entered")
-            return True
-    else:
-        enter_api_key_and_save_to_rc()
-        ph("cli-new-api-key-entered")
-        return True
 
+    enter_api_key_and_save_to_rc()
+    ph("cli-new-api-key-entered")
+    return True
 
-def enter_api_key_and_save_to_rc(existing_api_key: str = ""):
+
+def enter_api_key_and_save_to_rc():
     browser_launched = False
-    api_key = existing_api_key
+    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:
@@ -458,23 +526,36 @@ def enter_api_key_and_save_to_rc(existing_api_key: str = ""):
                 )
                 click.launch("https://app.codeflash.ai/app/apikeys")
                 browser_launched = True  # This does not work on remote consoles
-    shell_rc_path = os.path.expanduser(
-        f"~/.{os.environ.get('SHELL', '/bin/bash').split('/')[-1]}rc"
-    )
-    api_key_line = f'export CODEFLASH_API_KEY="{api_key}"'
-    api_key_pattern = re.compile(r'^export CODEFLASH_API_KEY=".*"$', re.M)
-    with open(shell_rc_path, "r+") as shell_rc:
-        shell_contents = shell_rc.read()
-        if api_key_pattern.search(shell_contents):
-            # Replace the existing API key line
-            updated_shell_contents = api_key_pattern.sub(api_key_line, shell_contents)
-        else:
-            # Append the new API key line
-            updated_shell_contents = shell_contents.rstrip() + f"\n{api_key_line}\n"
-        shell_rc.seek(0)
-        shell_rc.write(updated_shell_contents)
-        shell_rc.truncate()
-    click.echo(f"✅ Updated CODEFLASH_API_KEY in {shell_rc_path}")
+
+    shell_rc_path = get_shell_rc_path()
+    api_key_line = f"export CODEFLASH_API_KEY={api_key}"
+    try:
+        with open(shell_rc_path, "r+", encoding="utf8") as shell_file:
+            shell_contents = shell_file.read()
+            existing_api_key = read_api_key_from_shell_config()
+
+            if existing_api_key:
+                # Replace the existing API key line
+                updated_shell_contents = re.sub(
+                    SHELL_RC_EXPORT_PATTERN, api_key_line, shell_contents
+                )
+                action = "Updated CODEFLASH_API_KEY in"
+            else:
+                # Append the new API key line
+                updated_shell_contents = shell_contents.rstrip() + f"\n{api_key_line}\n"
+                action = "Added CODEFLASH_API_KEY to"
+
+            shell_file.seek(0)
+            shell_file.write(updated_shell_contents)
+            shell_file.truncate()
+        click.echo(f"✅ {action} {shell_rc_path}.")
+    except IOError as e:
+        click.echo(
+            f"💡 I tried adding your CodeFlash API key to {shell_rc_path} - but seems like I don't have permissions to do so.\n"
+            f"You'll need to open it yourself and add the following line:\n\n{api_key_line}\n"
+        )
+        click.pause()
+
     os.environ["CODEFLASH_API_KEY"] = api_key
 
 
@@ -489,7 +570,7 @@ def create_bubble_sort_file(setup_info: dict[str, str]):
     return arr
 """
     bubble_sort_path = os.path.join(setup_info["module_root"], "bubble_sort.py")
-    with open(bubble_sort_path, "w") as bubble_sort_file:
+    with open(bubble_sort_path, "w", encoding="utf8") as bubble_sort_file:
         bubble_sort_file.write(bubble_sort_content)
     click.echo(f"✅ Created {bubble_sort_path}")
 

{codeflash-0.3.4 → codeflash-0.3.6}/codeflash/cli_cmds/workflows/codeflash-optimize.yaml

@@ -19,26 +19,31 @@ jobs:
           echo "Checking if this PR is created by CodeFlash bot..."
           if [ "${{ github.event.pull_request.user.login }}" == "codeflash-ai[bot]" ]; then
             echo "PR created by CodeFlash bot. Skipping optimization."
-            exit 0
+            echo "skip_remaining_steps=yes" >> $GITHUB_OUTPUT
           else
+            echo "skip_remaining_steps=no" >> $GITHUB_OUTPUT
             echo "It's not. Proceeding with the optimization."
           fi
       - uses: actions/checkout@v3
+        if: steps.bot_check.outputs.skip_remaining_steps == 'no'
         with:
           fetch-depth: 0
           token: ${{ secrets.GITHUB_TOKEN }}
       # TODO: Replace this with your project's python installation method
       - name: Set up Python
+        if: steps.bot_check.outputs.skip_remaining_steps == 'no'
         uses: actions/setup-python@v4
         with:
           python-version: {{ python_version }}
       # TODO: Replace this with your project's dependency installation method
       - name: Install Project Dependencies
+        if: steps.bot_check.outputs.skip_remaining_steps == 'no'
         run: |
           python -m pip install --upgrade pip
           pip install -r requirements.txt # TODO: Replace this with your project setup method
           pip install codeflash
       - name: Run CodeFlash to optimize code
+        if: steps.bot_check.outputs.skip_remaining_steps == 'no'
         id: optimize_code
         run: |
           codeflash