Merge branch 'main' of https://github.com/codeflash-ai/codeflash into updated-vsc-extension

Author: KRRT7

codeflash/api/cfapi.py

@@ -45,13 +45,31 @@ def make_cfapi_request(
     cfapi_headers = {"Authorization": f"Bearer {get_codeflash_api_key()}"}
     if extra_headers:
         cfapi_headers.update(extra_headers)
-    if method.upper() == "POST":
-        json_payload = json.dumps(payload, indent=None, default=pydantic_encoder)
-        cfapi_headers["Content-Type"] = "application/json"
-        response = requests.post(url, data=json_payload, headers=cfapi_headers, timeout=60)
-    else:
-        response = requests.get(url, headers=cfapi_headers, timeout=60)
-    return response
+    try:
+        if method.upper() == "POST":
+            json_payload = json.dumps(payload, indent=None, default=pydantic_encoder)
+            cfapi_headers["Content-Type"] = "application/json"
+            response = requests.post(url, data=json_payload, headers=cfapi_headers, timeout=60)
+        else:
+            response = requests.get(url, headers=cfapi_headers, timeout=60)
+        response.raise_for_status()
+        return response  # noqa: TRY300
+    except requests.exceptions.HTTPError:
+        # response may be either a string or JSON, so we handle both cases
+        error_message = ""
+        try:
+            json_response = response.json()
+            if "error" in json_response:
+                error_message = json_response["error"]
+            elif "message" in json_response:
+                error_message = json_response["message"]
+        except (ValueError, TypeError):
+            error_message = response.text
+
+        logger.error(
+            f"CF_API_Error:: making request to Codeflash API (url: {url}, method: {method}, status {response.status_code}): {error_message}"
+        )
+        return response
 
 
 @lru_cache(maxsize=1)
@@ -165,10 +183,7 @@ def is_github_app_installed_on_repo(owner: str, repo: str) -> bool:
     :return: The response object.
     """
     response = make_cfapi_request(endpoint=f"/is-github-app-installed?repo={repo}&owner={owner}", method="GET")
-    if not response.ok or response.text != "true":
-        logger.error(f"Error: {response.text}")
-        return False
-    return True
+    return response.ok and response.text == "true"
 
 
 def get_blocklisted_functions() -> dict[str, set[str]] | dict[str, Any]:

codeflash/cli_cmds/cli.py

@@ -137,9 +137,6 @@ def process_pyproject_config(args: Namespace) -> Namespace:
         assert Path(args.benchmarks_root).is_dir(), (
             f"--benchmarks-root {args.benchmarks_root} must be a valid directory"
         )
-        assert Path(args.benchmarks_root).resolve().is_relative_to(Path(args.tests_root).resolve()), (
-            f"--benchmarks-root {args.benchmarks_root} must be a subdirectory of --tests-root {args.tests_root}"
-        )
         if env_utils.get_pr_number() is not None:
             import git
 

codeflash/code_utils/code_replacer.py

@@ -7,7 +7,7 @@
 
 import isort
 import libcst as cst
-import libcst.matchers as m
+from libcst.metadata import PositionProvider
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_extractor import add_global_assignments, add_needed_imports_from_module
@@ -37,6 +37,55 @@ def normalize_code(code: str) -> str:
     return ast.unparse(normalize_node(ast.parse(code)))
 
 
+class AddRequestArgument(cst.CSTTransformer):
+    METADATA_DEPENDENCIES = (PositionProvider,)
+
+    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
+        # Matcher for '@fixture' or '@pytest.fixture'
+        for decorator in original_node.decorators:
+            dec = decorator.decorator
+
+            if isinstance(dec, cst.Call):
+                func_name = ""
+                if isinstance(dec.func, cst.Attribute) and isinstance(dec.func.value, cst.Name):
+                    if dec.func.attr.value == "fixture" and dec.func.value.value == "pytest":
+                        func_name = "pytest.fixture"
+                elif isinstance(dec.func, cst.Name) and dec.func.value == "fixture":
+                    func_name = "fixture"
+
+                if func_name:
+                    for arg in dec.args:
+                        if (
+                            arg.keyword
+                            and arg.keyword.value == "autouse"
+                            and isinstance(arg.value, cst.Name)
+                            and arg.value.value == "True"
+                        ):
+                            args = updated_node.params.params
+                            arg_names = {arg.name.value for arg in args}
+
+                            # Skip if 'request' is already present
+                            if "request" in arg_names:
+                                return updated_node
+
+                            # Create a new 'request' param
+                            request_param = cst.Param(name=cst.Name("request"))
+
+                            # Add 'request' as the first argument (after 'self' or 'cls' if needed)
+                            if args:
+                                first_arg = args[0].name.value
+                                if first_arg in {"self", "cls"}:
+                                    new_params = [args[0], request_param] + list(args[1:])  # noqa: RUF005
+                                else:
+                                    new_params = [request_param] + list(args)  # noqa: RUF005
+                            else:
+                                new_params = [request_param]
+
+                            new_param_list = updated_node.params.with_changes(params=new_params)
+                            return updated_node.with_changes(params=new_param_list)
+        return updated_node
+
+
 class PytestMarkAdder(cst.CSTTransformer):
     """Transformer that adds pytest marks to test functions."""
 
@@ -106,41 +155,51 @@ def _create_pytest_mark(self) -> cst.Decorator:
 class AutouseFixtureModifier(cst.CSTTransformer):
     def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
         # Matcher for '@fixture' or '@pytest.fixture'
-        fixture_decorator_func = m.Name("fixture") | m.Attribute(value=m.Name("pytest"), attr=m.Name("fixture"))
-
         for decorator in original_node.decorators:
-            if m.matches(
-                decorator,
-                m.Decorator(
-                    decorator=m.Call(
-                        func=fixture_decorator_func, args=[m.Arg(value=m.Name("True"), keyword=m.Name("autouse"))]
-                    )
-                ),
-            ):
-                # Found a matching fixture with autouse=True
-
-                # 1. The original body of the function will become the 'else' block.
-                #    updated_node.body is an IndentedBlock, which is what cst.Else expects.
-                else_block = cst.Else(body=updated_node.body)
-
-                # 2. Create the new 'if' block that will exit the fixture early.
-                if_test = cst.parse_expression('request.node.get_closest_marker("codeflash_no_autouse")')
-                yield_statement = cst.parse_statement("yield")
-                if_body = cst.IndentedBlock(body=[yield_statement])
-
-                # 3. Construct the full if/else statement.
-                new_if_statement = cst.If(test=if_test, body=if_body, orelse=else_block)
-
-                # 4. Replace the entire function's body with our new single statement.
-                return updated_node.with_changes(body=cst.IndentedBlock(body=[new_if_statement]))
+            dec = decorator.decorator
+
+            if isinstance(dec, cst.Call):
+                func_name = ""
+                if isinstance(dec.func, cst.Attribute) and isinstance(dec.func.value, cst.Name):
+                    if dec.func.attr.value == "fixture" and dec.func.value.value == "pytest":
+                        func_name = "pytest.fixture"
+                elif isinstance(dec.func, cst.Name) and dec.func.value == "fixture":
+                    func_name = "fixture"
+
+                if func_name:
+                    for arg in dec.args:
+                        if (
+                            arg.keyword
+                            and arg.keyword.value == "autouse"
+                            and isinstance(arg.value, cst.Name)
+                            and arg.value.value == "True"
+                        ):
+                            # Found a matching fixture with autouse=True
+
+                            # 1. The original body of the function will become the 'else' block.
+                            #    updated_node.body is an IndentedBlock, which is what cst.Else expects.
+                            else_block = cst.Else(body=updated_node.body)
+
+                            # 2. Create the new 'if' block that will exit the fixture early.
+                            if_test = cst.parse_expression('request.node.get_closest_marker("codeflash_no_autouse")')
+                            yield_statement = cst.parse_statement("yield")
+                            if_body = cst.IndentedBlock(body=[yield_statement])
+
+                            # 3. Construct the full if/else statement.
+                            new_if_statement = cst.If(test=if_test, body=if_body, orelse=else_block)
+
+                            # 4. Replace the entire function's body with our new single statement.
+                            return updated_node.with_changes(body=cst.IndentedBlock(body=[new_if_statement]))
         return updated_node
 
 
 def disable_autouse(test_path: Path) -> str:
     file_content = test_path.read_text(encoding="utf-8")
     module = cst.parse_module(file_content)
+    add_request_argument = AddRequestArgument()
     disable_autouse_fixture = AutouseFixtureModifier()
-    modified_module = module.visit(disable_autouse_fixture)
+    modified_module = module.visit(add_request_argument)
+    modified_module = modified_module.visit(disable_autouse_fixture)
     test_path.write_text(modified_module.code, encoding="utf-8")
     return file_content
 

codeflash/tracer.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import contextlib
+import datetime
 import importlib.machinery
 import io
 import json
@@ -81,6 +82,7 @@ def __init__(
         config_file_path: Path | None = None,
         max_function_count: int = 256,
         timeout: int | None = None,  # seconds
+        command: str | None = None,
     ) -> None:
         """Use this class to trace function calls.
 
@@ -91,6 +93,7 @@ def __init__(
         :param max_function_count: Maximum number of times to trace one function
         :param timeout: Timeout in seconds for the tracer, if the traced code takes more than this time, then tracing
                     stops and normal execution continues. If this is None then no timeout applies
+        :param command: The command that initiated the tracing (for metadata storage)
         """
         if functions is None:
             functions = []
@@ -148,6 +151,9 @@ def __init__(
         assert "test_framework" in self.config, "Please specify 'test-framework' in pyproject.toml config file"
         self.t = self.timer()
 
+        # Store command information for metadata table
+        self.command = command if command else " ".join(sys.argv)
+
     def __enter__(self) -> None:
         if self.disable:
             return
@@ -174,6 +180,22 @@ def __enter__(self) -> None:
             "CREATE TABLE function_calls(type TEXT, function TEXT, classname TEXT, filename TEXT, "
             "line_number INTEGER, last_frame_address INTEGER, time_ns INTEGER, args BLOB)"
         )
+
+        # Create metadata table to store command information
+        cur.execute("CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT)")
+
+        # Store command metadata
+        cur.execute("INSERT INTO metadata VALUES (?, ?)", ("command", self.command))
+        cur.execute("INSERT INTO metadata VALUES (?, ?)", ("program_name", self.file_being_called_from))
+        cur.execute(
+            "INSERT INTO metadata VALUES (?, ?)",
+            ("functions_filter", json.dumps(self.functions) if self.functions else None),
+        )
+        cur.execute(
+            "INSERT INTO metadata VALUES (?, ?)",
+            ("timestamp", datetime.datetime.now(datetime.timezone.utc).isoformat()),
+        )
+        cur.execute("INSERT INTO metadata VALUES (?, ?)", ("project_root", str(self.project_root)))
         console.rule("Codeflash: Traced Program Output Begin", style="bold blue")
         frame = sys._getframe(0)  # Get this frame and simulate a call to it  # noqa: SLF001
         self.dispatch["call"](self, frame, 0)
@@ -842,6 +864,7 @@ def main() -> ArgumentParser:
                 max_function_count=args.max_function_count,
                 timeout=args.tracer_timeout,
                 config_file_path=args.codeflash_config,
+                command=" ".join(sys.argv),
             ).runctx(code, globs, None)
 
         except BrokenPipeError as exc:

codeflash/version.py

@@ -1,3 +1,3 @@
 # These version placeholders will be replaced by poetry-dynamic-versioning during `poetry build`.
-__version__ = "0.14.0"
-__version_tuple__ = (0, 14, 0)
+__version__ = "0.14.3"
+__version_tuple__ = (0, 14, 3)

docs/docs/configuration.md

@@ -16,6 +16,7 @@ tests-root = "tests"
 test-framework = "pytest"
 formatter-cmds = ["black $file"]
 # optional configuration
+benchmarks-root = "tests/benchmarks" # Required when running with --benchmark
 ignore-paths = ["my_module/build/"]
 pytest-cmd = "pytest"
 disable-imports-sorting = false
@@ -29,6 +30,7 @@ Required Options:
 - `test-framework`: The test framework you use for your project. Codeflash supports `pytest` and `unittest`.
 
 Optional Configuration:
+- `benchmarks-root`: The directory where your benchmarks are located. Codeflash will use this directory to discover existing benchmarks. Note that this option is required when running with `--benchmark`.
 - `ignore-paths`: A list of paths withing the `module-root` to ignore when optimizing code. Codeflash will not optimize code in these paths. Useful for ignoring build directories or other generated code. You can also leave this empty if not needed.
 - `pytest-cmd`: The command to run your tests. Defaults to `pytest`. You can specify extra commandline arguments here for pytest.
 - `formatter-cmds`: The command line to run your code formatter or linter. Defaults to `["black $file"]`. In the command line `$file` refers to the current file being optimized. The assumption with using tools here is that they overwrite the same file and returns a zero exit code. You can also specify multiple tools here that run in a chain as a toml array. You can also disable code formatting by setting this to `["disabled"]`.

docs/docs/optimizing-with-codeflash/benchmarking.md

@@ -3,7 +3,7 @@ sidebar_position: 5
 ---
 # Using Benchmarks
 
-Codeflash is able the determine the impact of an optimization on predefined benchmarks, when used in benchmark mode.
+Codeflash is able to determine the impact of an optimization on predefined benchmarks, when used in benchmark mode.
 
 Benchmark mode is an easy way for users to define workflows that are performance-critical and need to be optimized.
 For example, if a user has an important function that requires minimal latency, the user can define a benchmark for that function.
@@ -13,7 +13,7 @@ Codeflash will then calculate the impact (if any) of any optimization on the per
 
 1. **Create a benchmarks root** 
 
-    Create a directory for benchmarks. This directory must be a sub directory of your tests directory.
+    Create a directory for benchmarks if it does not already exist. 
 
    In your pyproject.toml, add the path to the 'benchmarks-root' section.
     ```yaml