Merge remote-tracking branch 'origin/main' into stdout_comparison_

Author: KRRT7

codeflash/LICENSE

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

codeflash/code_utils/code_replacer.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import ast
+import re
 from collections import defaultdict
 from functools import lru_cache
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Optional, TypeVar
 
 import libcst as cst
 

codeflash/code_utils/concolic_utils.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import ast
+import re
+from typing import Optional
+
+
+class AssertCleanup:
+    def transform_asserts(self, code: str) -> str:
+        lines = code.splitlines()
+        result_lines = []
+
+        for line in lines:
+            transformed = self._transform_assert_line(line)
+            result_lines.append(transformed if transformed is not None else line)
+
+        return "\n".join(result_lines)
+
+    def _transform_assert_line(self, line: str) -> Optional[str]:
+        indent = line[: len(line) - len(line.lstrip())]
+
+        assert_match = self.assert_re.match(line)
+        if assert_match:
+            expression = assert_match.group(1).strip()
+            if expression.startswith("not "):
+                return f"{indent}{expression}"
+
+            expression = expression.rstrip(",;")
+            return f"{indent}{expression}"
+
+        unittest_match = self.unittest_re.match(line)
+        if unittest_match:
+            indent, assert_method, args = unittest_match.groups()
+
+            if args:
+                arg_parts = self._split_top_level_args(args)
+                if arg_parts and arg_parts[0]:
+                    return f"{indent}{arg_parts[0]}"
+
+        return None
+
+    def _split_top_level_args(self, args_str: str) -> list[str]:
+        result = []
+        current = []
+        depth = 0
+
+        for char in args_str:
+            if char in "([{":
+                depth += 1
+                current.append(char)
+            elif char in ")]}":
+                depth -= 1
+                current.append(char)
+            elif char == "," and depth == 0:
+                result.append("".join(current).strip())
+                current = []
+            else:
+                current.append(char)
+
+        if current:
+            result.append("".join(current).strip())
+
+        return result
+
+    def __init__(self):
+        # Pre-compiling regular expressions for faster execution
+        self.assert_re = re.compile(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$")
+        self.unittest_re = re.compile(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$")
+
+
+def clean_concolic_tests(test_suite_code: str) -> str:
+    try:
+        can_parse = True
+        tree = ast.parse(test_suite_code)
+    except SyntaxError:
+        can_parse = False
+
+    if not can_parse:
+        return AssertCleanup().transform_asserts(test_suite_code)
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"):
+            new_body = []
+            for stmt in node.body:
+                if isinstance(stmt, ast.Assert):
+                    if isinstance(stmt.test, ast.Compare) and isinstance(stmt.test.left, ast.Call):
+                        new_body.append(ast.Expr(value=stmt.test.left))
+                    else:
+                        new_body.append(stmt)
+
+                else:
+                    new_body.append(stmt)
+            node.body = new_body
+
+    return ast.unparse(tree).strip()

codeflash/discovery/functions_to_optimize.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import json
 import os
 import random
 import warnings
@@ -156,9 +157,9 @@ def get_functions_to_optimize(
     project_root: Path,
     module_root: Path,
 ) -> tuple[dict[Path, list[FunctionToOptimize]], int]:
-    assert (
-        sum([bool(optimize_all), bool(replay_test), bool(file)]) <= 1
-    ), "Only one of optimize_all, replay_test, or file should be provided"
+    assert sum([bool(optimize_all), bool(replay_test), bool(file)]) <= 1, (
+        "Only one of optimize_all, replay_test, or file should be provided"
+    )
     functions: dict[str, list[FunctionToOptimize]]
     with warnings.catch_warnings():
         warnings.simplefilter(action="ignore", category=SyntaxWarning)
@@ -434,9 +435,7 @@ def filter_functions(
             test_functions_removed_count += len(functions)
             continue
         if file_path in ignore_paths or any(
-            # file_path.startswith(ignore_path + os.sep) for ignore_path in ignore_paths if ignore_path
-            file_path.startswith(str(ignore_path) + os.sep)
-            for ignore_path in ignore_paths
+            file_path.startswith(str(ignore_path) + os.sep) for ignore_path in ignore_paths
         ):
             ignore_paths_removed_count += 1
             continue
@@ -457,15 +456,17 @@ def filter_functions(
             malformed_paths_count += 1
             continue
         if blocklist_funcs:
-            for function in functions.copy():
-                path = Path(function.file_path).name
-                if path in blocklist_funcs and function.function_name in blocklist_funcs[path]:
-                    functions.remove(function)
-                    logger.debug(f"Skipping {function.function_name} in {path} as it has already been optimized")
-                    continue
-
+            functions = [
+                function
+                for function in functions
+                if not (
+                    function.file_path.name in blocklist_funcs
+                    and function.qualified_name in blocklist_funcs[function.file_path.name]
+                )
+            ]
         filtered_modified_functions[file_path] = functions
         functions_count += len(functions)
+
     if not disable_logs:
         log_info = {
             f"{test_functions_removed_count} test function{'s' if test_functions_removed_count != 1 else ''}": test_functions_removed_count,
@@ -475,10 +476,11 @@ def filter_functions(
             f"{ignore_paths_removed_count} file{'s' if ignore_paths_removed_count != 1 else ''} from ignored paths": ignore_paths_removed_count,
             f"{submodule_ignored_paths_count} file{'s' if submodule_ignored_paths_count != 1 else ''} from ignored submodules": submodule_ignored_paths_count,
         }
-        log_string: str
-        if log_string := "\n".join([k for k, v in log_info.items() if v > 0]):
+        log_string = "\n".join([k for k, v in log_info.items() if v > 0])
+        if log_string:
             logger.info(f"Ignoring: {log_string}")
             console.rule()
+
     return {Path(k): v for k, v in filtered_modified_functions.items() if v}, functions_count
 
 

codeflash/github/PrComment.py

@@ -20,6 +20,13 @@ class PrComment:
     winning_benchmarking_test_results: TestResults
 
     def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str]]:
+
+        report_table = {
+            test_type.to_name(): result
+            for test_type, result in self.winning_behavioral_test_results.get_test_pass_fail_report_by_type().items()
+            if test_type.to_name()
+        }
+
         return {
             "optimization_explanation": self.optimization_explanation,
             "best_runtime": humanize_runtime(self.best_runtime),
@@ -29,10 +36,7 @@ def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str]]:
             "speedup_x": self.speedup_x,
             "speedup_pct": self.speedup_pct,
             "loop_count": self.winning_benchmarking_test_results.number_of_loops(),
-            "report_table": {
-                test_type.to_name(): result
-                for test_type, result in self.winning_behavioral_test_results.get_test_pass_fail_report_by_type().items()
-            },
+            "report_table": report_table
         }
 
 

codeflash/models/models.py

@@ -7,12 +7,13 @@
 from enum import Enum, IntEnum
 from pathlib import Path
 from re import Pattern
-from typing import Any, Optional, Union
+from typing import Annotated, Any, Optional, Union
 
+import sentry_sdk
+from coverage.exceptions import NoDataError
 from jedi.api.classes import Name
 from pydantic import AfterValidator, BaseModel, ConfigDict, Field
 from pydantic.dataclasses import dataclass
-from typing_extensions import Annotated
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.code_utils import validate_python_code
@@ -217,7 +218,7 @@ class CoverageData:
     graph: dict[str, dict[str, Collection[object]]]
     code_context: CodeOptimizationContext
     main_func_coverage: FunctionCoverage
-    dependent_func_coverage: Union[FunctionCoverage, None]
+    dependent_func_coverage: Optional[FunctionCoverage]
     status: CoverageStatus
     blank_re: Pattern[str] = re.compile(r"\s*(#|$)")
     else_re: Pattern[str] = re.compile(r"\s*else\s*:\s*(#|$)")
@@ -231,34 +232,21 @@ def load_from_sqlite_database(
         from coverage.jsonreport import JsonReporter
 
         cov = Coverage(data_file=database_path, data_suffix=True, auto_data=True, branch=True)
+
         if not database_path.stat().st_size or not database_path.exists():
             logger.debug(f"Coverage database {database_path} is empty or does not exist")
-            return CoverageData(
-                file_path=source_code_path,
-                coverage=0.0,
-                function_name=function_name,
-                functions_being_tested=[],
-                graph={},
-                code_context=code_context,
-                main_func_coverage=FunctionCoverage(
-                    name=function_name,
-                    coverage=0.0,
-                    executed_lines=[],
-                    unexecuted_lines=[],
-                    executed_branches=[],
-                    unexecuted_branches=[],
-                ),
-                dependent_func_coverage=None,
-                status=CoverageStatus.NOT_FOUND,
-            )
-
+            sentry_sdk.capture_message(f"Coverage database {database_path} is empty or does not exist")
+            return CoverageData.create_empty(source_code_path, function_name, code_context)
         cov.load()
 
         reporter = JsonReporter(cov)
         temp_json_file = database_path.with_suffix(".report.json")
         with temp_json_file.open("w") as f:
-            reporter.report(morfs=[source_code_path.as_posix()], outfile=f)
-
+            try:
+                reporter.report(morfs=[source_code_path.as_posix()], outfile=f)
+            except NoDataError:
+                sentry_sdk.capture_message(f"No coverage data found for {function_name} in {source_code_path}")
+                return CoverageData.create_empty(source_code_path, function_name, code_context)
         with temp_json_file.open() as f:
             original_coverage_data = json.load(f)
 
@@ -461,6 +449,34 @@ def log_coverage(self) -> None:
         if is_end_to_end():
             console.print(self)
 
+    @classmethod
+    def create_empty(cls, file_path: Path, function_name: str, code_context: CodeOptimizationContext) -> CoverageData:
+        return cls(
+            file_path=file_path,
+            coverage=0.0,
+            function_name=function_name,
+            functions_being_tested=[function_name],
+            graph={
+                function_name: {
+                    "executed_lines": set(),
+                    "unexecuted_lines": set(),
+                    "executed_branches": [],
+                    "unexecuted_branches": [],
+                }
+            },
+            code_context=code_context,
+            main_func_coverage=FunctionCoverage(
+                name=function_name,
+                coverage=0.0,
+                executed_lines=[],
+                unexecuted_lines=[],
+                executed_branches=[],
+                unexecuted_branches=[],
+            ),
+            dependent_func_coverage=None,
+            status=CoverageStatus.NOT_FOUND,
+        )
+
 
 @dataclass
 class FunctionCoverage:

codeflash/verification/concolic_testing.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 
 from codeflash.cli_cmds.console import console, logger
+from codeflash.code_utils.concolic_utils import clean_concolic_tests
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
 from codeflash.code_utils.static_analysis import has_typed_parameters
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
@@ -21,7 +22,11 @@ def generate_concolic_tests(
 ) -> tuple[dict[str, list[FunctionCalledInTest]], str]:
     function_to_concolic_tests = {}
     concolic_test_suite_code = ""
-    if test_cfg.concolic_test_root_dir and has_typed_parameters(function_to_optimize_ast, function_to_optimize.parents):
+    if (
+        test_cfg.concolic_test_root_dir
+        and isinstance(function_to_optimize_ast, (ast.FunctionDef, ast.AsyncFunctionDef))
+        and has_typed_parameters(function_to_optimize_ast, function_to_optimize.parents)
+    ):
         logger.info("Generating concolic opcode coverage tests for the original code…")
         console.rule()
         try:
@@ -54,7 +59,8 @@ def generate_concolic_tests(
             return function_to_concolic_tests, concolic_test_suite_code
 
         if cover_result.returncode == 0:
-            concolic_test_suite_code: str = cover_result.stdout
+            generated_concolic_test: str = cover_result.stdout
+            concolic_test_suite_code: str = clean_concolic_tests(generated_concolic_test)
             concolic_test_suite_dir = Path(tempfile.mkdtemp(dir=test_cfg.concolic_test_root_dir))
             concolic_test_suite_path = concolic_test_suite_dir / "test_concolic_coverage.py"
             concolic_test_suite_path.write_text(concolic_test_suite_code, encoding="utf8")

codeflash/version.py

@@ -1,3 +1,3 @@
 # These version placeholders will be replaced by poetry-dynamic-versioning during `poetry build`.
-__version__ = "0.9.2"
-__version_tuple__ = (0, 9, 2)
+__version__ = "0.10.0"
+__version_tuple__ = (0, 10, 0)

pyproject.toml

@@ -67,7 +67,7 @@ exclude = [
 
 # Versions here the minimum required versions for the project. These should be as loose as possible.
 [tool.poetry.dependencies]
-python = "^3.9"
+python = ">=3.9"
 unidiff = ">=0.7.4"
 pytest = ">=7.0.0"
 gitpython = ">=3.1.31"
@@ -176,7 +176,7 @@ ignore = [
     "TD003",
     "TD004",
     "PLR2004",
-    "UP007" # remove once we drop 3.9 support.
+    "UP007"
 ]
 
 [tool.ruff.lint.flake8-type-checking]