codeflash 0.3.3__tar.gz → 0.3.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: codeflash
-Version: 0.3.3
+Version: 0.3.4
 Summary: A Code Performance Optimization Library
 Home-page: https://codeflash.ai
 Author: CodeFlash Inc.
@@ -20,7 +20,7 @@ Requires-Dist: jedi (>=0.19.0)
 Requires-Dist: junitparser (>=3.1.0)
 Requires-Dist: libcst (>=1.0.1)
 Requires-Dist: posthog (>=3.0.0)
-Requires-Dist: pydantic (>=2.5.2,<3.0.0)
+Requires-Dist: pydantic (>=1.10.1)
 Requires-Dist: pytest (>=7.0.0)
 Requires-Dist: pytest-timeout (>=2.1.0)
 Requires-Dist: tiktoken (>=0.3.2)

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/api/aiservice.py

@@ -1,9 +1,11 @@
+import json
 import logging
 import os
-import requests
-from pydantic import RootModel
 from typing import Any, Dict, List, Tuple, Optional
 
+import requests
+from pydantic.json import pydantic_encoder
+
 from codeflash.analytics.posthog import ph
 from codeflash.code_utils.env_utils import get_codeflash_api_key
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
@@ -23,19 +25,20 @@ def make_ai_service_request(
 ) -> requests.Response:
     """
     Make an API request to the given endpoint on the AI service.
-
-    Parameters:
-    - endpoint (str): The endpoint to call, e.g., "/optimize".
-    - method (str): The HTTP method to use, e.g., "POST".
-    - data (Dict[str, Any]): The data to send in the request.
-
-    Returns:
-    - requests.Response: The response from the API.
+    :param endpoint: The endpoint to call, e.g., "/optimize".
+    :param method: The HTTP method to use ('GET' or 'POST').
+    :param payload: Optional JSON payload to include in the POST request body.
+    :param timeout: The timeout for the request.
+    :return: The response object from the API.
     """
     url = f"{AI_SERVICE_BASE_URL}/ai{endpoint}"
     ai_service_headers = {"Authorization": f"Bearer {get_codeflash_api_key()}"}
     if method.upper() == "POST":
-        response = requests.post(url, json=payload, headers=ai_service_headers, timeout=timeout)
+        json_payload = json.dumps(payload, indent=None, default=pydantic_encoder)
+        ai_service_headers["Content-Type"] = "application/json"
+        response = requests.post(
+            url, data=json_payload, headers=ai_service_headers, timeout=timeout
+        )
     else:
         response = requests.get(url, headers=ai_service_headers, timeout=timeout)
     # response.raise_for_status()  # Will raise an HTTPError if the HTTP request returned an unsuccessful status code
@@ -55,10 +58,10 @@ def optimize_python_code(
     Returns:
     - List[Tuple[str, str]]: A list of tuples where the first element is the optimized code and the second is the explanation.
     """
-    data = {"source_code": source_code, "num_variants": num_variants}
+    payload = {"source_code": source_code, "num_variants": num_variants}
     logging.info(f"Generating optimized candidates ...")
     try:
-        response = make_ai_service_request("/optimize", payload=data, timeout=600)
+        response = make_ai_service_request("/optimize", payload=payload, timeout=600)
     except requests.exceptions.RequestException as e:
         logging.error(f"Error generating optimized candidates: {e}")
         ph("cli-optimize-error-caught", {"error": str(e)})
@@ -66,6 +69,7 @@ def optimize_python_code(
 
     if response.status_code == 200:
         optimizations = response.json()["optimizations"]
+        logging.info(f"Generated {len(optimizations)} candidates.")
         return [(opt["source_code"], opt["explanation"]) for opt in optimizations]
     else:
         try:
@@ -108,11 +112,9 @@ def generate_regression_tests(
         "pytest",
         "unittest",
     ], f"Invalid test framework, got {test_framework} but expected 'pytest' or 'unittest'"
-    data = {
+    payload = {
         "source_code_being_tested": source_code_being_tested,
-        "function_to_optimize": RootModel[FunctionToOptimize](function_to_optimize).model_dump(
-            mode="json"
-        ),
+        "function_to_optimize": function_to_optimize,
         "dependent_function_names": dependent_function_names,
         "module_path": module_path,
         "test_module_path": test_module_path,
@@ -120,7 +122,7 @@ def generate_regression_tests(
         "test_timeout": test_timeout,
     }
     try:
-        response = make_ai_service_request("/testgen", payload=data, timeout=600)
+        response = make_ai_service_request("/testgen", payload=payload, timeout=600)
     except requests.exceptions.RequestException as e:
         logging.error(f"Error generating tests: {e}")
         ph("cli-testgen-error-caught", {"error": str(e)})
@@ -130,6 +132,7 @@ def generate_regression_tests(
 
     if response.status_code == 200:
         response_json = response.json()
+        logging.info(f"Generated tests for function {function_to_optimize.function_name}")
         return response_json["generated_tests"], response_json["instrumented_tests"]
     else:
         try:

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

@@ -1,9 +1,11 @@
+import json
 import logging
 import os
 from functools import lru_cache
 from typing import Optional, Dict, Any
 
 import requests
+from pydantic.json import pydantic_encoder
 from requests import Response
 
 from codeflash.code_utils.env_utils import get_codeflash_api_key
@@ -21,15 +23,17 @@ def make_cfapi_request(
 ) -> requests.Response:
     """
     Make an HTTP request using the specified method, URL, headers, and JSON payload.
-    :param endpoint: The URL to send the request to.
+    :param endpoint: The endpoint URL to send the request to.
     :param method: The HTTP method to use ('GET', 'POST', etc.).
-    :param payload: Optional JSON payload to include in the request body.
-    :return: The response object.
+    :param payload: Optional JSON payload to include in the POST request body.
+    :return: The response object from the API.
     """
     url = f"{CFAPI_BASE_URL}/cfapi{endpoint}"
     cfapi_headers = {"Authorization": f"Bearer {get_codeflash_api_key()}"}
     if method.upper() == "POST":
-        response = requests.post(url, json=payload, headers=cfapi_headers)
+        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)
     else:
         response = requests.get(url, headers=cfapi_headers)
     return response

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/code_utils/env_utils.py

@@ -1,8 +1,10 @@
 import logging
 import os
+from functools import lru_cache
 from typing import Optional
 
 
+@lru_cache(maxsize=1)
 def get_codeflash_api_key() -> Optional[str]:
     api_key = os.environ.get("CODEFLASH_API_KEY")
     if not api_key:
@@ -33,11 +35,13 @@ def ensure_codeflash_api_key() -> bool:
     return True
 
 
+@lru_cache(maxsize=1)
 def get_codeflash_org_key() -> Optional[str]:
     api_key = os.environ.get("CODEFLASH_ORG_KEY")
     return api_key
 
 
+@lru_cache(maxsize=1)
 def get_pr_number() -> Optional[int]:
     pr_number = os.environ.get("CODEFLASH_PR_NUMBER")
     if not pr_number:

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/discovery/functions_to_optimize.py

@@ -1,12 +1,12 @@
 import ast
+import libcst as cst
 import logging
 import os
+import random
 from _ast import ClassDef, FunctionDef, AsyncFunctionDef
-from typing import Dict, Optional, List, Tuple, Union
-
-import libcst as cst
 from libcst import CSTNode
 from pydantic.dataclasses import dataclass
+from typing import Dict, Optional, List, Tuple, Union
 
 from codeflash.code_utils.code_utils import path_belongs_to_site_packages
 from codeflash.code_utils.git_utils import get_git_diff
@@ -128,7 +128,10 @@ def get_functions_to_optimize_by_file(
                 if only_function_name == fn.function_name:
                     found_function = fn
             if found_function is None:
-                raise ValueError(f"Function {only_function_name} not found in file {file}")
+                raise ValueError(
+                    f"Function {only_function_name} not found in file {file} or"
+                    f" the function does not have a 'return' statement."
+                )
             functions[file] = [found_function]
     else:
         logging.info("Finding all functions modified in the current git diff ...")
@@ -170,7 +173,12 @@ def get_all_files_and_functions(module_root_path: str) -> Dict[str, List[Functio
             file_path = os.path.join(root, file)
             # Find all the functions in the file
             functions.update(find_all_functions_in_file(file_path))
-    return functions
+    # Randomize the order of the files to optimize to avoid optimizing the same file in the same order every time.
+    # Helpful if an optimize-all run is stuck and we restart it.
+    files_list = list(functions.items())
+    random.shuffle(files_list)
+    functions_shuffled = dict(files_list)
+    return functions_shuffled
 
 
 def find_all_functions_in_file(file_path: str) -> Dict[str, List[FunctionToOptimize]]:

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/main.py

@@ -1,3 +1,6 @@
+"""Thanks for being curious about how codeflash works! If you might want to work with us on finally making performance a
+solved problem, please reach out to us at careers@codeflash.ai. We're hiring!
+"""
 import concurrent.futures
 import libcst as cst
 import logging

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/optimization/function_context.py

@@ -1,12 +1,12 @@
 import ast
-import jedi
 import logging
 import os
+from typing import List
+
+import jedi
 import tiktoken
 from jedi.api.classes import Name
-from pydantic import RootModel
 from pydantic.dataclasses import dataclass
-from typing import List
 
 from codeflash.code_utils.code_extractor import get_code_no_skeleton, get_code
 from codeflash.code_utils.code_utils import path_belongs_to_site_packages
@@ -43,11 +43,6 @@ class Source:
     source_code: str
 
 
-@dataclass(frozen=True)
-class SourceList(RootModel):
-    root: list[Source]
-
-
 def get_type_annotation_context(
     function: FunctionToOptimize, jedi_script: jedi.Script, project_root_path: str
 ) -> List[Source]:
@@ -72,12 +67,18 @@ def get_type_annotation_context(
                             name = arg.annotation.id
                             line_no = arg.annotation.lineno
                             col_no = arg.annotation.col_offset
-                            definition: List[Name] = jedi_script.goto(
-                                line=line_no,
-                                column=col_no,
-                                follow_imports=True,
-                                follow_builtin_imports=False,
-                            )
+                            try:
+                                definition: List[Name] = jedi_script.goto(
+                                    line=line_no,
+                                    column=col_no,
+                                    follow_imports=True,
+                                    follow_builtin_imports=False,
+                                )
+                            except Exception as e:
+                                logging.error(
+                                    f"Error while getting definition for {name.full_name}: {e}"
+                                )
+                                definition = []
                             if definition:  # TODO can be multiple definitions
                                 definition_path = str(definition[0].module_path)
                                 # The definition is part of this project and not defined within the original function
@@ -136,12 +137,16 @@ def get_function_variables_definitions(
                     names.append(ref)
 
     for name in names:
-        definitions: List[Name] = script.goto(
-            line=name.line,
-            column=name.column,
-            follow_imports=True,
-            follow_builtin_imports=False,
-        )
+        try:
+            definitions: List[Name] = script.goto(
+                line=name.line,
+                column=name.column,
+                follow_imports=True,
+                follow_builtin_imports=False,
+            )
+        except Exception as e:
+            logging.error(f"Error while getting definition for {name.full_name}: {e}")
+            definitions = []
         if definitions:
             # TODO: there can be multiple definitions, see how to handle such cases
             definition_path = str(definitions[0].module_path)

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/result/create_pr.py

@@ -29,9 +29,7 @@ def check_create_pr(
             repo=repo,
             pr_number=pr_number,
             file_changes={
-                relative_path: FileDiffContent(
-                    oldContent=original_code, newContent=new_code
-                ).model_dump(mode="json")
+                relative_path: FileDiffContent(oldContent=original_code, newContent=new_code)
             },
             pr_comment=PrComment(
                 optimization_explanation=explanation.explanation_message(),
@@ -64,9 +62,7 @@ def check_create_pr(
             repo=repo,
             base_branch=base_branch,
             file_changes={
-                relative_path: FileDiffContent(
-                    oldContent=original_code, newContent=new_code
-                ).model_dump(mode="json")
+                relative_path: FileDiffContent(oldContent=original_code, newContent=new_code)
             },
             pr_comment=PrComment(
                 optimization_explanation=explanation.explanation_message(),

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/verification/comparator.py

@@ -1,4 +1,5 @@
 import datetime
+import decimal
 import logging
 import math
 from typing import Any
@@ -15,6 +16,19 @@ try:
     HAS_SQLALCHEMY = True
 except ImportError:
     HAS_SQLALCHEMY = False
+try:
+    import scipy
+
+    HAS_SCIPY = True
+except ImportError:
+    HAS_SCIPY = False
+
+try:
+    import pandas
+
+    HAS_PANDAS = True
+except ImportError:
+    HAS_PANDAS = False
 
 
 def comparator(orig: Any, new: Any) -> bool:
@@ -43,13 +57,14 @@ def comparator(orig: Any, new: Any) -> bool:
                 return False
         return True
 
-    if isinstance(orig, (str, int, bool, complex, type(None))):
+    if isinstance(orig, (str, int, bool, complex, type(None), decimal.Decimal)):
         return orig == new
     if isinstance(orig, float):
         if math.isnan(orig) and math.isnan(new):
             return True
         return math.isclose(orig, new)
-    if isinstance(orig, dict):
+    # scipy condition because dok_matrix type is also a instance of dict, but dict comparison doesn't work for it
+    if isinstance(orig, dict) and not (HAS_SCIPY and isinstance(orig, scipy.sparse.spmatrix)):
         if len(orig) != len(new):
             return False
         for key in orig:
@@ -76,6 +91,28 @@ def comparator(orig: Any, new: Any) -> bool:
     if HAS_NUMPY and isinstance(orig, (np.integer, np.bool_, np.byte)):
         return orig == new
 
+    if HAS_SCIPY and isinstance(orig, scipy.sparse.spmatrix):
+        if orig.dtype != new.dtype:
+            return False
+        if orig.get_shape() != new.get_shape():
+            return False
+        return (orig != new).nnz == 0
+
+    if HAS_PANDAS and isinstance(
+        orig,
+        (
+            pandas.DataFrame,
+            pandas.Series,
+            pandas.Index,
+            pandas.Categorical,
+            pandas.arrays.SparseArray,
+        ),
+    ):
+        return orig.equals(new)
+
+    if HAS_PANDAS and isinstance(orig, (pandas.CategoricalDtype, pandas.Interval, pandas.Period)):
+        return orig == new
+
     # This should be at the end of all numpy checking
     try:
         if HAS_NUMPY and np.isnan(orig):

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/verification/parse_test_output.py

@@ -1,3 +1,5 @@
+from collections import defaultdict
+
 import logging
 import os
 import pathlib
@@ -5,10 +7,8 @@ import pickle
 import re
 import sqlite3
 import subprocess
-from collections import defaultdict
-from typing import Optional
-
 from junitparser.xunit2 import JUnitXml
+from typing import Optional
 
 from codeflash.code_utils.code_utils import (
     module_name_from_file_path,
@@ -258,6 +258,7 @@ def merge_test_results(xml_test_results: TestResults, bin_test_results: TestResu
                     )
                 )
         else:
+            # Looks like we would never reach this condition because we are not printing the log line in the xml file
             for xml_result in xml_results:
                 bin_result = bin_results.get_by_id(xml_result.id)
                 if bin_result is None:

{codeflash-0.3.3 → codeflash-0.3.4}/codeflash/version.py

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

{codeflash-0.3.3 → codeflash-0.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "codeflash"
-version = "0.3.3" # Determined by poetry-dynamic-versioning during `poetry build`
+version = "0.3.4" # Determined by poetry-dynamic-versioning during `poetry build`
 description = "A Code Performance Optimization Library"
 authors = ["CodeFlash Inc. <contact@codeflash.ai>"]
 homepage = "https://codeflash.ai"
@@ -22,7 +22,7 @@ pytest-timeout = ">=2.1.0"
 tomlkit = ">=0.11.7"
 unittest-xml-reporting = ">=3.2.0"
 junitparser = ">=3.1.0"
-pydantic = "^2.5.2"
+pydantic = ">=1.10.1"
 black = ">=22.3.0"
 humanize = ">=4.0.0"
 posthog = ">=3.0.0"
@@ -60,7 +60,7 @@ line-length = 100
 target-version = ['py312']
 
 [build-system]
-requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]
+requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.2.0,<2.0.0"]
 build-backend = "poetry_dynamic_versioning.backend"
 
 [tool.poetry.build]