Merge branch 'main' into alpha-async

Author: KRRT7

README.md

@@ -3,9 +3,7 @@
    <a href="https://github.com/codeflash-ai/codeflash">
     <img src="https://img.shields.io/github/commit-activity/m/codeflash-ai/codeflash" alt="GitHub commit activity">
   </a>
-  <a href="https://pypi.org/project/codeflash/">
-    <img src="https://img.shields.io/pypi/dm/codeflash" alt="PyPI Downloads">
-  </a>
+  <a href="https://pypi.org/project/codeflash/"><img src="https://static.pepy.tech/badge/codeflash" alt="PyPI Downloads"></a>
   <a href="https://pypi.org/project/codeflash/">
     <img src="https://img.shields.io/pypi/v/codeflash?label=PyPI%20version" alt="PyPI Downloads">
   </a>

codeflash/code_utils/code_extractor.py

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

codeflash/code_utils/formatter.py

@@ -44,19 +44,14 @@ def apply_formatter_cmds(
     test_dir_str: Optional[str],
     print_status: bool,  # noqa
     exit_on_failure: bool = True,  # noqa
-) -> tuple[Path, str]:
-    # TODO: Only allow a particular whitelist of formatters here to prevent arbitrary code execution
-    formatter_name = cmds[0].lower()
+) -> tuple[Path, str, bool]:
     should_make_copy = False
     file_path = path
 
     if test_dir_str:
         should_make_copy = True
         file_path = Path(test_dir_str) / "temp.py"
 
-    if not cmds or formatter_name == "disabled":
-        return path, path.read_text(encoding="utf8")
-
     if not path.exists():
         msg = f"File {path} does not exist. Cannot apply formatter commands."
         raise FileNotFoundError(msg)
@@ -66,6 +61,7 @@ def apply_formatter_cmds(
 
     file_token = "$file"  # noqa: S105
 
+    changed = False
     for command in cmds:
         formatter_cmd_list = shlex.split(command, posix=os.name != "nt")
         formatter_cmd_list = [file_path.as_posix() if chunk == file_token else chunk for chunk in formatter_cmd_list]
@@ -74,6 +70,7 @@ def apply_formatter_cmds(
             if result.returncode == 0:
                 if print_status:
                     console.rule(f"Formatted Successfully with: {command.replace('$file', path.name)}")
+                changed = True
             else:
                 logger.error(f"Failed to format code with {' '.join(formatter_cmd_list)}")
         except FileNotFoundError as e:
@@ -88,7 +85,7 @@ def apply_formatter_cmds(
             if exit_on_failure:
                 raise e from None
 
-    return file_path, file_path.read_text(encoding="utf8")
+    return file_path, file_path.read_text(encoding="utf8"), changed
 
 
 def get_diff_lines_count(diff_output: str) -> int:
@@ -112,10 +109,16 @@ def format_code(
     if console.quiet:
         # lsp mode
         exit_on_failure = False
-    with tempfile.TemporaryDirectory() as test_dir_str:
-        if isinstance(path, str):
-            path = Path(path)
 
+    if isinstance(path, str):
+        path = Path(path)
+
+    # TODO: Only allow a particular whitelist of formatters here to prevent arbitrary code execution
+    formatter_name = formatter_cmds[0].lower() if formatter_cmds else "disabled"
+    if formatter_name == "disabled":
+        return path.read_text(encoding="utf8")
+
+    with tempfile.TemporaryDirectory() as test_dir_str:
         original_code = path.read_text(encoding="utf8")
         original_code_lines = len(original_code.split("\n"))
 
@@ -126,26 +129,39 @@ def format_code(
             original_temp = Path(test_dir_str) / "original_temp.py"
             original_temp.write_text(original_code_without_opfunc, encoding="utf8")
 
-            formatted_temp, formatted_code = apply_formatter_cmds(
-                formatter_cmds, original_temp, test_dir_str, print_status=False
+            formatted_temp, formatted_code, changed = apply_formatter_cmds(
+                formatter_cmds, original_temp, test_dir_str, print_status=False, exit_on_failure=exit_on_failure
             )
 
+            if not changed:
+                logger.warning(
+                    f"No changes detected in {path} after formatting, are you sure you have valid formatter commands?"
+                )
+                return original_code
+
             diff_output = generate_unified_diff(
                 original_code_without_opfunc, formatted_code, from_file=str(original_temp), to_file=str(formatted_temp)
             )
             diff_lines_count = get_diff_lines_count(diff_output)
 
             max_diff_lines = min(int(original_code_lines * 0.3), 50)
 
-            if diff_lines_count > max_diff_lines and max_diff_lines != -1:
-                logger.debug(
+            if diff_lines_count > max_diff_lines:
+                logger.warning(
                     f"Skipping formatting {path}: {diff_lines_count} lines would change (max: {max_diff_lines})"
                 )
                 return original_code
+
         # TODO : We can avoid formatting the whole file again and only formatting the optimized code standalone and replace in formatted file above.
-        _, formatted_code = apply_formatter_cmds(
+        _, formatted_code, changed = apply_formatter_cmds(
             formatter_cmds, path, test_dir_str=None, print_status=print_status, exit_on_failure=exit_on_failure
         )
+        if not changed:
+            logger.warning(
+                f"No changes detected in {path} after formatting, are you sure you have valid formatter commands?"
+            )
+            return original_code
+
         logger.debug(f"Formatted {path} with commands: {formatter_cmds}")
         return formatted_code
 

codeflash/optimization/function_optimizer.py

@@ -739,7 +739,9 @@ def reformat_code_and_helpers(
             file_to_code_context = optimized_context.file_to_path()
             optimized_code = file_to_code_context.get(str(path.relative_to(self.project_root)), "")
 
-        new_code = format_code(self.args.formatter_cmds, path, optimized_code=optimized_code, check_diff=True)
+        new_code = format_code(
+            self.args.formatter_cmds, path, optimized_code=optimized_code, check_diff=True, exit_on_failure=False
+        )
         if should_sort_imports:
             new_code = sort_imports(new_code)
 
@@ -748,7 +750,11 @@ def reformat_code_and_helpers(
             module_abspath = hp.file_path
             hp_source_code = hp.source_code
             formatted_helper_code = format_code(
-                self.args.formatter_cmds, module_abspath, optimized_code=hp_source_code, check_diff=True
+                self.args.formatter_cmds,
+                module_abspath,
+                optimized_code=hp_source_code,
+                check_diff=True,
+                exit_on_failure=False,
             )
             if should_sort_imports:
                 formatted_helper_code = sort_imports(formatted_helper_code)

codeflash/version.py

@@ -1,2 +1,2 @@
 # These version placeholders will be replaced by uv-dynamic-versioning during build.
-__version__ = "0.16.4"
+__version__ = "0.16.5"