Merge branch 'main' into saga4/misc_lsp_12_sept

Author: Saga4

codeflash/code_utils/code_extractor.py

@@ -335,12 +335,12 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
         return updated_node
 
 
-def extract_global_statements(source_code: str) -> list[cst.SimpleStatementLine]:
+def extract_global_statements(source_code: str) -> tuple[cst.Module, list[cst.SimpleStatementLine]]:
     """Extract global statements from source code."""
     module = cst.parse_module(source_code)
     collector = GlobalStatementCollector()
     module.visit(collector)
-    return collector.global_statements
+    return module, collector.global_statements
 
 
 def find_last_import_line(target_code: str) -> int:
@@ -373,30 +373,41 @@ def delete___future___aliased_imports(module_code: str) -> str:
 
 
 def add_global_assignments(src_module_code: str, dst_module_code: str) -> str:
-    non_assignment_global_statements = extract_global_statements(src_module_code)
+    src_module, new_added_global_statements = extract_global_statements(src_module_code)
+    dst_module, existing_global_statements = extract_global_statements(dst_module_code)
 
-    # Find the last import line in target
-    last_import_line = find_last_import_line(dst_module_code)
-
-    # Parse the target code
-    target_module = cst.parse_module(dst_module_code)
-
-    # Create transformer to insert non_assignment_global_statements
-    transformer = ImportInserter(non_assignment_global_statements, last_import_line)
-    #
-    # # Apply transformation
-    modified_module = target_module.visit(transformer)
-    dst_module_code = modified_module.code
-
-    # Parse the code
-    original_module = cst.parse_module(dst_module_code)
-    new_module = cst.parse_module(src_module_code)
+    unique_global_statements = []
+    for stmt in new_added_global_statements:
+        if any(
+            stmt is existing_stmt or stmt.deep_equals(existing_stmt) for existing_stmt in existing_global_statements
+        ):
+            continue
+        unique_global_statements.append(stmt)
+
+    mod_dst_code = dst_module_code
+    # Insert unique global statements if any
+    if unique_global_statements:
+        last_import_line = find_last_import_line(dst_module_code)
+        # Reuse already-parsed dst_module
+        transformer = ImportInserter(unique_global_statements, last_import_line)
+        # Use visit inplace, don't parse again
+        modified_module = dst_module.visit(transformer)
+        mod_dst_code = modified_module.code
+        # Parse the code after insertion
+        original_module = cst.parse_module(mod_dst_code)
+    else:
+        # No new statements to insert, reuse already-parsed dst_module
+        original_module = dst_module
 
+    # Parse the src_module_code once only (already done above: src_module)
     # Collect assignments from the new file
     new_collector = GlobalAssignmentCollector()
-    new_module.visit(new_collector)
+    src_module.visit(new_collector)
+    # Only create transformer if there are assignments to insert/transform
+    if not new_collector.assignments:  # nothing to transform
+        return mod_dst_code
 
-    # Transform the original file
+    # Transform the original destination module
     transformer = GlobalAssignmentTransformer(new_collector.assignments, new_collector.assignment_order)
     transformed_module = original_module.visit(transformer)
 

codeflash/code_utils/code_replacer.py

@@ -412,11 +412,17 @@ def replace_function_definitions_in_module(
     module_abspath: Path,
     preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
     project_root_path: Path,
+    should_add_global_assignments: bool = True,  # noqa: FBT001, FBT002
 ) -> bool:
     source_code: str = module_abspath.read_text(encoding="utf8")
     code_to_apply = get_optimized_code_for_module(module_abspath.relative_to(project_root_path), optimized_code)
+
     new_code: str = replace_functions_and_add_imports(
-        add_global_assignments(code_to_apply, source_code),
+        # adding the global assignments before replacing the code, not after
+        # becuase of an "edge case" where the optimized code intoduced a new import and a global assignment using that import
+        # and that import wasn't used before, so it was ignored when calling AddImportsVisitor.add_needed_import inside replace_functions_and_add_imports (because the global assignment wasn't added yet)
+        # this was added at https://github.com/codeflash-ai/codeflash/pull/448
+        add_global_assignments(code_to_apply, source_code) if should_add_global_assignments else source_code,
         function_names,
         code_to_apply,
         module_abspath,

codeflash/code_utils/shell_utils.py

@@ -15,7 +15,9 @@
     SHELL_RC_EXPORT_PATTERN = re.compile(r"^set CODEFLASH_API_KEY=(cf-.*)$", re.MULTILINE)
     SHELL_RC_EXPORT_PREFIX = "set CODEFLASH_API_KEY="
 else:
-    SHELL_RC_EXPORT_PATTERN = re.compile(r'^(?!#)export CODEFLASH_API_KEY=[\'"]?(cf-[^\s"]+)[\'"]$', re.MULTILINE)
+    SHELL_RC_EXPORT_PATTERN = re.compile(
+        r'^(?!#)export CODEFLASH_API_KEY=(?:"|\')?(cf-[^\s"\']+)(?:"|\')?$', re.MULTILINE
+    )
     SHELL_RC_EXPORT_PREFIX = "export CODEFLASH_API_KEY="
 
 

codeflash/code_utils/version_check.py

@@ -0,0 +1,79 @@
+"""Version checking utilities for codeflash."""
+
+from __future__ import annotations
+
+import time
+
+import requests
+from packaging import version
+
+from codeflash.cli_cmds.console import console, logger
+from codeflash.version import __version__
+
+# Simple cache to avoid checking too frequently
+_version_cache = {"version": '0.0.0', "timestamp": float(0)}
+_cache_duration = 3600  # 1 hour cache
+
+
+def get_latest_version_from_pypi() -> str | None:
+    """Get the latest version of codeflash from PyPI.
+
+    Returns:
+        The latest version string from PyPI, or None if the request fails.
+
+    """
+    # Check cache first
+    current_time = time.time()
+    if _version_cache["version"] is not None and current_time - _version_cache["timestamp"] < _cache_duration:
+        return _version_cache["version"]
+
+    try:
+        response = requests.get("https://pypi.org/pypi/codeflash/json", timeout=2)
+        if response.status_code == 200:
+            data = response.json()
+            latest_version = data["info"]["version"]
+
+            # Update cache
+            _version_cache["version"] = latest_version
+            _version_cache["timestamp"] = current_time
+
+            return latest_version
+        logger.debug(f"Failed to fetch version from PyPI: {response.status_code}")
+        return None  # noqa: TRY300
+    except requests.RequestException as e:
+        logger.debug(f"Network error fetching version from PyPI: {e}")
+        return None
+    except (KeyError, ValueError) as e:
+        logger.debug(f"Invalid response format from PyPI: {e}")
+        return None
+    except Exception as e:
+        logger.debug(f"Unexpected error fetching version from PyPI: {e}")
+        return None
+
+
+def check_for_newer_minor_version() -> None:
+    """Check if a newer minor version is available on PyPI and notify the user.
+
+    This function compares the current version with the latest version on PyPI.
+    If a newer minor version is available, it prints an informational message
+    suggesting the user upgrade.
+    """
+    latest_version = get_latest_version_from_pypi()
+
+    if not latest_version:
+        return
+
+    try:
+        current_parsed = version.parse(__version__)
+        latest_parsed = version.parse(latest_version)
+
+        # Check if there's a newer minor version available
+        # We only notify for minor version updates, not patch updates
+        if latest_parsed > current_parsed: # < > == operators can be directly applied on version objects
+            logger.warning(
+                f"A newer version({latest_version}) of Codeflash is available, please update soon!"
+            )
+
+    except version.InvalidVersion as e:
+        logger.debug(f"Invalid version format: {e}")
+        return

codeflash/context/unused_definition_remover.py

@@ -537,6 +537,7 @@ def revert_unused_helper_functions(
                     module_abspath=file_path,
                     preexisting_objects=set(),  # Empty set since we're reverting
                     project_root_path=project_root,
+                    should_add_global_assignments=False,  # since we revert helpers functions after applying the optimization, we know that the file already has global assignments added, otherwise they would be added twice.
                 )
 
                 if reverted_code:

codeflash/main.py

@@ -11,6 +11,7 @@
 from codeflash.cli_cmds.console import paneled_text
 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
 
@@ -21,12 +22,15 @@ def main() -> None:
         CODEFLASH_LOGO, panel_args={"title": "https://codeflash.ai", "expand": False}, text_args={"style": "bold gold3"}
     )
     args = parse_args()
+
+    # Check for newer version for all commands
+    check_for_newer_minor_version()
+
     if args.command:
+        disable_telemetry = False
         if args.config_file and Path.exists(args.config_file):
             pyproject_config, _ = parse_config_file(args.config_file)
             disable_telemetry = pyproject_config.get("disable_telemetry", False)
-        else:
-            disable_telemetry = False
         init_sentry(not disable_telemetry, exclude_errors=True)
         posthog_cf.initialize_posthog(not disable_telemetry)
         args.func()

codeflash/optimization/function_optimizer.py

@@ -53,6 +53,7 @@
 )
 from codeflash.code_utils.env_utils import get_pr_number
 from codeflash.code_utils.formatter import format_code, sort_imports
+from codeflash.code_utils.git_utils import git_root_dir
 from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test
 from codeflash.code_utils.line_profile_utils import add_decorator_imports
 from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast
@@ -822,7 +823,10 @@ def reformat_code_and_helpers(
         return new_code, new_helper_code
 
     def replace_function_and_helpers_with_optimized_code(
-        self, code_context: CodeOptimizationContext, optimized_code: CodeStringsMarkdown, original_helper_code: str
+        self,
+        code_context: CodeOptimizationContext,
+        optimized_code: CodeStringsMarkdown,
+        original_helper_code: dict[Path, str],
     ) -> bool:
         did_update = False
         read_writable_functions_by_file_path = defaultdict(set)
@@ -1302,11 +1306,13 @@ def process_review(
             "coverage_message": coverage_message,
             "replay_tests": replay_tests,
             "concolic_tests": concolic_tests,
-            "root_dir": self.project_root,
         }
 
         raise_pr = not self.args.no_pr
 
+        if raise_pr or self.args.staging_review:
+            data["root_dir"] = git_root_dir()
+
         if raise_pr and not self.args.staging_review:
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)