initial test files

Author: mohammedahmed18

code_to_optimize/code_directories/circular_deps/api_client.py

@@ -0,0 +1,77 @@
+from os import getenv
+from typing import Optional
+
+from attrs import define, evolve, field
+
+from code_to_optimize.code_directories.circular_deps.constants import DEFAULT_API_URL, DEFAULT_APP_URL
+
+
+@define
+class GalileoApiClient():
+    """A Client which has been authenticated for use on secured endpoints
+
+    The following are accepted as keyword arguments and will be used to construct httpx Clients internally:
+
+        ``base_url``: The base URL for the API, all requests are made to a relative path to this URL
+            This can also be set via the GALILEO_CONSOLE_URL environment variable
+
+        ``api_key``: The API key to be sent with every request
+            This can also be set via the GALILEO_API_KEY environment variable
+
+        ``cookies``: A dictionary of cookies to be sent with every request
+
+        ``headers``: A dictionary of headers to be sent with every request
+
+        ``timeout``: The maximum amount of a time a request can take. API functions will raise
+        httpx.TimeoutException if this is exceeded.
+
+        ``verify_ssl``: Whether or not to verify the SSL certificate of the API server. This should be True in production,
+        but can be set to False for testing purposes.
+
+        ``follow_redirects``: Whether or not to follow redirects. Default value is False.
+
+        ``httpx_args``: A dictionary of additional arguments to be passed to the ``httpx.Client`` and ``httpx.AsyncClient`` constructor.
+
+    Attributes:
+        raise_on_unexpected_status: Whether or not to raise an errors.UnexpectedStatus if the API returns a
+            status code that was not documented in the source OpenAPI document. Can also be provided as a keyword
+            argument to the constructor.
+        token: The token to use for authentication
+        prefix: The prefix to use for the Authorization header
+        auth_header_name: The name of the Authorization header
+    """
+
+    _base_url: Optional[str] = field(factory=lambda: GalileoApiClient.get_api_url(), kw_only=True, alias="base_url")
+    _api_key: Optional[str] = field(factory=lambda: getenv("GALILEO_API_KEY", None), kw_only=True, alias="api_key")
+    token: Optional[str] = None
+
+    api_key_header_name: str = "Galileo-API-Key"
+    client_type_header_name: str = "client-type"
+    client_type_header_value: str = "sdk-python"
+
+    @staticmethod
+    def get_console_url() -> str:
+        console_url = getenv("GALILEO_CONSOLE_URL", DEFAULT_API_URL)
+        if DEFAULT_API_URL == console_url:
+            return DEFAULT_APP_URL
+
+        return console_url
+
+    def with_api_key(self, api_key: str) -> "GalileoApiClient":
+        """Get a new client matching this one with a new API key"""
+        if self._client is not None:
+            self._client.headers.update({self.api_key_header_name: api_key})
+        if self._async_client is not None:
+            self._async_client.headers.update({self.api_key_header_name: api_key})
+        return evolve(self, api_key=api_key)
+
+    @staticmethod
+    def get_api_url(base_url: Optional[str] = None) -> str:
+        api_url = base_url or getenv("GALILEO_CONSOLE_URL", DEFAULT_API_URL)
+        if api_url is None:
+            raise ValueError("base_url or GALILEO_CONSOLE_URL must be set")
+        if any(map(api_url.__contains__, ["localhost", "127.0.0.1"])):
+            api_url = "http://localhost:8088"
+        else:
+            api_url = api_url.replace("app.galileo.ai", "api.galileo.ai").replace("console", "api")
+        return api_url

code_to_optimize/code_directories/circular_deps/constants.py

@@ -0,0 +1,8 @@
+DEFAULT_API_URL = "https://api.galileo.ai/"
+DEFAULT_APP_URL = "https://app.galileo.ai/"
+
+
+# function_names: GalileoApiClient.get_console_url
+# module_abs_path : /home/mohammed/Work/galileo-python/src/galileo/api_client.py
+# preexisting_objects: {('GalileoApiClient', ()), ('_set_destination', ()), ('get_console_url', (FunctionParent(name='GalileoApiClient', type='ClassDef'),))}
+# project_root_path: /home/mohammed/Work/galileo-python/src

codeflash/code_utils/code_replacer.py

@@ -397,6 +397,13 @@ def replace_functions_and_add_imports(
     preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
     project_root_path: Path,
 ) -> str:
+    logger.debug("start from here,...")
+    logger.debug(f"source_code: {source_code}")
+    logger.debug(f"function_names: {function_names}")
+    logger.debug(f"optimized_code: {optimized_code}")
+    logger.debug(f"module_abspath: {module_abspath}")
+    logger.debug(f"preexisting_objects: {preexisting_objects}")
+    logger.debug(f"project_root_path: {project_root_path}")
     return add_needed_imports_from_module(
         optimized_code,
         replace_functions_in_file(source_code, function_names, optimized_code, preexisting_objects),

tests/test_code_context_extractor.py

@@ -2434,3 +2434,53 @@ def simple_method(self):
         assert "class SimpleClass:" in code_content
         assert "def simple_method(self):" in code_content
         assert "return 42" in code_content
+
+
+from codeflash.code_utils.code_replacer import replace_functions_and_add_imports
+def test_replace_functions_and_add_imports():
+    path_to_root = Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "circular_deps"
+    optimized_code = '''from __future__ import annotations
+
+import urllib.parse
+from os import getenv
+
+from attrs import define
+from code_to_optimize.code_directories.circular_deps.constants import DEFAULT_API_URL, DEFAULT_APP_URL
+
+# Precompute constant netlocs for set membership test
+_DEFAULT_APP_NETLOC = urllib.parse.urlparse(DEFAULT_APP_URL).netloc
+_DEFAULT_API_NETLOC = urllib.parse.urlparse(DEFAULT_API_URL).netloc
+_NETLOC_SET = {_DEFAULT_APP_NETLOC, _DEFAULT_API_NETLOC}
+
+@define
+class GalileoApiClient():
+
+    @staticmethod
+    def get_console_url() -> str:
+        # Return DEFAULT_APP_URL if the env var is not set or set to DEFAULT_API_URL
+        console_url = getenv("GALILEO_CONSOLE_URL", DEFAULT_API_URL)
+        if console_url == DEFAULT_API_URL:
+            return DEFAULT_APP_URL
+        return console_url
+
+def _set_destination(console_url: str) -> str:
+    """
+    Parse the console_url and return the destination for the OpenTelemetry traces.
+    """
+    destination = (console_url or GalileoApiClient.get_console_url()).replace("console.", "api.")
+    parsed_url = urllib.parse.urlparse(destination)
+    if parsed_url.netloc in _NETLOC_SET:
+        return f"{DEFAULT_APP_URL}api/galileo/otel/traces"
+    return f"{parsed_url.scheme}://{parsed_url.netloc}/otel/traces"'''
+    file_abs_path = path_to_root / "api_client.py"
+    content = Path(file_abs_path).read_text(encoding="utf-8")
+    new_code = replace_functions_and_add_imports(
+        source_code= content,
+        function_names= ["GalileoApiClient.get_console_url"],
+        optimized_code= optimized_code,
+        module_abspath= file_abs_path,
+        preexisting_objects= {('GalileoApiClient', ()), ('_set_destination', ()), ('get_console_url', (FunctionParent(name='GalileoApiClient', type='ClassDef'),))},
+        project_root_path= Path(path_to_root),
+    )
+    print(new_code)
+    assert 1 == 1