codeflash 0.7.2__tar.gz → 0.7.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (68) hide show
  1. {codeflash-0.7.2 → codeflash-0.7.4}/PKG-INFO +2 -1
  2. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/api/aiservice.py +14 -27
  3. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/api/cfapi.py +34 -32
  4. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/cli_cmds/cli.py +8 -24
  5. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/cli_cmds/cli_common.py +4 -18
  6. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/cli_cmds/cmd_init.py +32 -94
  7. codeflash-0.7.4/codeflash/cli_cmds/console.py +67 -0
  8. codeflash-0.7.4/codeflash/cli_cmds/console_constants.py +65 -0
  9. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/cli_cmds/logging_config.py +3 -16
  10. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/code_extractor.py +14 -43
  11. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/code_replacer.py +17 -58
  12. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/code_utils.py +13 -3
  13. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/config_parser.py +3 -10
  14. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/env_utils.py +5 -5
  15. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/formatter.py +2 -10
  16. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/git_utils.py +5 -21
  17. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/github_utils.py +2 -2
  18. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/instrument_existing_tests.py +62 -168
  19. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/shell_utils.py +10 -19
  20. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/time_utils.py +1 -4
  21. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/discovery/discover_unit_tests.py +31 -47
  22. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/discovery/functions_to_optimize.py +19 -55
  23. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/main.py +1 -3
  24. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/models/models.py +10 -21
  25. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/optimization/function_context.py +20 -64
  26. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/optimization/optimizer.py +264 -388
  27. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/result/create_pr.py +7 -10
  28. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/result/critic.py +24 -31
  29. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/telemetry/posthog_cf.py +2 -9
  30. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/tracer.py +21 -86
  31. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/tracing/profile_stats.py +2 -6
  32. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/tracing/replay_test.py +13 -32
  33. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/comparator.py +18 -31
  34. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/parse_test_output.py +63 -83
  35. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/pytest_plugin.py +7 -25
  36. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/test_results.py +17 -26
  37. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/test_runner.py +1 -3
  38. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/verification_utils.py +5 -8
  39. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/verifier.py +7 -26
  40. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/version.py +2 -2
  41. {codeflash-0.7.2 → codeflash-0.7.4}/pyproject.toml +28 -16
  42. codeflash-0.7.2/codeflash/cli_cmds/console.py +0 -44
  43. {codeflash-0.7.2 → codeflash-0.7.4}/README.md +0 -0
  44. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/.env.example +0 -0
  45. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/LICENSE +0 -0
  46. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/__init__.py +0 -0
  47. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/api/__init__.py +0 -0
  48. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/cli_cmds/__init__.py +0 -0
  49. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/cli_cmds/workflows/codeflash-optimize.yaml +0 -0
  50. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/__init__.py +0 -0
  51. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/compat.py +0 -0
  52. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/config_consts.py +0 -0
  53. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/code_utils/remove_generated_tests.py +0 -0
  54. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/discovery/__init__.py +0 -0
  55. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/github/PrComment.py +0 -0
  56. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/github/__init__.py +0 -0
  57. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/models/ExperimentMetadata.py +0 -0
  58. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/models/__init__.py +0 -0
  59. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/optimization/__init__.py +0 -0
  60. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/result/__init__.py +0 -0
  61. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/result/explanation.py +0 -0
  62. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/telemetry/__init__.py +0 -0
  63. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/telemetry/sentry.py +0 -0
  64. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/tracing/__init__.py +0 -0
  65. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/tracing/tracing_utils.py +0 -0
  66. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/update_license_version.py +0 -0
  67. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/__init__.py +0 -0
  68. {codeflash-0.7.2 → codeflash-0.7.4}/codeflash/verification/equivalence.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: codeflash
3
- Version: 0.7.2
3
+ Version: 0.7.4
4
4
  Summary: Client for codeflash.ai - automatic code performance optimization, powered by AI
5
5
  Home-page: https://codeflash.ai
6
6
  License: BSL-1.1
@@ -24,6 +24,7 @@ Requires-Dist: isort (>=5.11.0)
24
24
  Requires-Dist: jedi (>=0.19.1)
25
25
  Requires-Dist: junitparser (>=3.1.0)
26
26
  Requires-Dist: libcst (>=1.0.1)
27
+ Requires-Dist: lxml (>=5.3.0,<6.0.0)
27
28
  Requires-Dist: parameterized (>=0.9.0)
28
29
  Requires-Dist: posthog (>=3.0.0)
29
30
  Requires-Dist: pydantic (>=1.10.1)
@@ -9,7 +9,7 @@ import requests
9
9
  from pydantic.dataclasses import dataclass
10
10
  from pydantic.json import pydantic_encoder
11
11
 
12
- from codeflash.cli_cmds.console import logger
12
+ from codeflash.cli_cmds.console import console, logger
13
13
  from codeflash.code_utils.env_utils import get_codeflash_api_key
14
14
  from codeflash.discovery.functions_to_optimize import FunctionToOptimize
15
15
  from codeflash.models.ExperimentMetadata import ExperimentMetadata
@@ -33,10 +33,7 @@ class OptimizedCandidate:
33
33
  class AiServiceClient:
34
34
  def __init__(self) -> None:
35
35
  self.base_url = self.get_aiservice_base_url()
36
- self.headers = {
37
- "Authorization": f"Bearer {get_codeflash_api_key()}",
38
- "Connection": "close",
39
- }
36
+ self.headers = {"Authorization": f"Bearer {get_codeflash_api_key()}", "Connection": "close"}
40
37
 
41
38
  def get_aiservice_base_url(self) -> str:
42
39
  if os.environ.get("CODEFLASH_AIS_SERVER", default="prod").lower() == "local":
@@ -45,11 +42,7 @@ class AiServiceClient:
45
42
  return "https://app.codeflash.ai"
46
43
 
47
44
  def make_ai_service_request(
48
- self,
49
- endpoint: str,
50
- method: str = "POST",
51
- payload: dict[str, Any] | None = None,
52
- timeout: float | None = None,
45
+ self, endpoint: str, method: str = "POST", payload: dict[str, Any] | None = None, timeout: float | None = None
53
46
  ) -> requests.Response:
54
47
  """Make an API request to the given endpoint on the AI service.
55
48
 
@@ -96,13 +89,11 @@ class AiServiceClient:
96
89
  "experiment_metadata": experiment_metadata,
97
90
  "codeflash_version": codeflash_version,
98
91
  }
92
+
99
93
  logger.info("Generating optimized candidates ...")
94
+ console.rule()
100
95
  try:
101
- response = self.make_ai_service_request(
102
- "/optimize",
103
- payload=payload,
104
- timeout=600,
105
- )
96
+ response = self.make_ai_service_request("/optimize", payload=payload, timeout=600)
106
97
  except requests.exceptions.RequestException as e:
107
98
  logger.exception(f"Error generating optimized candidates: {e}")
108
99
  ph("cli-optimize-error-caught", {"error": str(e)})
@@ -111,6 +102,7 @@ class AiServiceClient:
111
102
  if response.status_code == 200:
112
103
  optimizations_json = response.json()["optimizations"]
113
104
  logger.info(f"Generated {len(optimizations_json)} candidates.")
105
+ console.rule()
114
106
  return [
115
107
  OptimizedCandidate(
116
108
  source_code=opt["source_code"],
@@ -124,10 +116,8 @@ class AiServiceClient:
124
116
  except Exception:
125
117
  error = response.text
126
118
  logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
127
- ph(
128
- "cli-optimize-error-response",
129
- {"response_status_code": response.status_code, "error": error},
130
- )
119
+ ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
120
+ console.rule()
131
121
  return []
132
122
 
133
123
  def log_results(
@@ -225,20 +215,17 @@ class AiServiceClient:
225
215
  try:
226
216
  error = response.json()["error"]
227
217
  logger.error(f"Error generating tests: {response.status_code} - {error}")
228
- ph(
229
- "cli-testgen-error-response",
230
- {"response_status_code": response.status_code, "error": error},
231
- )
218
+ ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": error})
232
219
  return None
233
220
  except Exception:
234
221
  logger.error(f"Error generating tests: {response.status_code} - {response.text}")
235
- ph(
236
- "cli-testgen-error-response",
237
- {"response_status_code": response.status_code, "error": response.text},
238
- )
222
+ ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": response.text})
239
223
  return None
240
224
 
241
225
 
242
226
  class LocalAiServiceClient(AiServiceClient):
227
+ """Client for interacting with the local AI service."""
228
+
243
229
  def get_aiservice_base_url(self) -> str:
230
+ """Get the base URL for the local AI service."""
244
231
  return "http://localhost:8000"
@@ -4,17 +4,19 @@ import json
4
4
  import os
5
5
  from functools import lru_cache
6
6
  from pathlib import Path
7
- from typing import Any, Dict, Optional
7
+ from typing import TYPE_CHECKING, Any, Dict, Optional
8
8
 
9
9
  import requests
10
10
  from pydantic.json import pydantic_encoder
11
- from requests import Response
12
11
 
13
12
  from codeflash.cli_cmds.console import logger
14
13
  from codeflash.code_utils.env_utils import ensure_codeflash_api_key, get_codeflash_api_key, get_pr_number
15
14
  from codeflash.code_utils.git_utils import get_repo_owner_and_name
16
15
  from codeflash.github.PrComment import FileDiffContent, PrComment
17
16
 
17
+ if TYPE_CHECKING:
18
+ from requests import Response
19
+
18
20
  if os.environ.get("CODEFLASH_CFAPI_SERVER", default="prod").lower() == "local":
19
21
  CFAPI_BASE_URL = "http://localhost:3001"
20
22
  logger.info(f"Using local CF API at {CFAPI_BASE_URL}.")
@@ -22,12 +24,9 @@ else:
22
24
  CFAPI_BASE_URL = "https://app.codeflash.ai"
23
25
 
24
26
 
25
- def make_cfapi_request(
26
- endpoint: str,
27
- method: str,
28
- payload: Optional[Dict[str, Any]] = None,
29
- ) -> requests.Response:
27
+ def make_cfapi_request(endpoint: str, method: str, payload: dict[str, Any] | None = None) -> Response:
30
28
  """Make an HTTP request using the specified method, URL, headers, and JSON payload.
29
+
31
30
  :param endpoint: The endpoint URL to send the request to.
32
31
  :param method: The HTTP method to use ('GET', 'POST', etc.).
33
32
  :param payload: Optional JSON payload to include in the POST request body.
@@ -38,15 +37,16 @@ def make_cfapi_request(
38
37
  if method.upper() == "POST":
39
38
  json_payload = json.dumps(payload, indent=None, default=pydantic_encoder)
40
39
  cfapi_headers["Content-Type"] = "application/json"
41
- response = requests.post(url, data=json_payload, headers=cfapi_headers)
40
+ response = requests.post(url, data=json_payload, headers=cfapi_headers, timeout=60)
42
41
  else:
43
- response = requests.get(url, headers=cfapi_headers)
42
+ response = requests.get(url, headers=cfapi_headers, timeout=60)
44
43
  return response
45
44
 
46
45
 
47
46
  @lru_cache(maxsize=1)
48
47
  def get_user_id() -> Optional[str]:
49
48
  """Retrieve the user's userid by making a request to the /cfapi/cli-get-user endpoint.
49
+
50
50
  :return: The userid or None if the request fails.
51
51
  """
52
52
  if not ensure_codeflash_api_key():
@@ -55,11 +55,8 @@ def get_user_id() -> Optional[str]:
55
55
  response = make_cfapi_request(endpoint="/cli-get-user", method="GET")
56
56
  if response.status_code == 200:
57
57
  return response.text
58
- else:
59
- logger.error(
60
- f"Failed to look up your userid; is your CF API key valid? ({response.reason})",
61
- )
62
- return None
58
+ logger.error(f"Failed to look up your userid; is your CF API key valid? ({response.reason})")
59
+ return None
63
60
 
64
61
 
65
62
  def suggest_changes(
@@ -73,6 +70,7 @@ def suggest_changes(
73
70
  trace_id: str,
74
71
  ) -> Response:
75
72
  """Suggest changes to a pull request.
73
+
76
74
  Will make a review suggestion when possible;
77
75
  or create a new dependent pull request with the suggested changes.
78
76
  :param owner: The owner of the repository.
@@ -93,8 +91,7 @@ def suggest_changes(
93
91
  "generatedTests": generated_tests,
94
92
  "traceId": trace_id,
95
93
  }
96
- response = make_cfapi_request(endpoint="/suggest-pr-changes", method="POST", payload=payload)
97
- return response
94
+ return make_cfapi_request(endpoint="/suggest-pr-changes", method="POST", payload=payload)
98
95
 
99
96
 
100
97
  def create_pr(
@@ -107,7 +104,8 @@ def create_pr(
107
104
  generated_tests: str,
108
105
  trace_id: str,
109
106
  ) -> Response:
110
- """Create a pull request, targeting the specified branch. (usually 'main')
107
+ """Create a pull request, targeting the specified branch. (usually 'main').
108
+
111
109
  :param owner: The owner of the repository.
112
110
  :param repo: The name of the repository.
113
111
  :param base_branch: The base branch to target.
@@ -133,37 +131,41 @@ def create_pr(
133
131
 
134
132
  def is_github_app_installed_on_repo(owner: str, repo: str) -> bool:
135
133
  """Check if the Codeflash GitHub App is installed on the specified repository.
134
+
136
135
  :param owner: The owner of the repository.
137
136
  :param repo: The name of the repository.
138
137
  :return: The response object.
139
138
  """
140
- response = make_cfapi_request(
141
- endpoint=f"/is-github-app-installed?repo={repo}&owner={owner}",
142
- method="GET",
143
- )
139
+ response = make_cfapi_request(endpoint=f"/is-github-app-installed?repo={repo}&owner={owner}", method="GET")
144
140
  if not response.ok or response.text != "true":
145
141
  logger.error(f"Error: {response.text}")
146
142
  return False
147
143
  return True
148
144
 
149
145
 
150
- def get_blocklisted_functions() -> dict[str, str]:
146
+ def get_blocklisted_functions() -> dict[str, set[str]]:
147
+ """Retrieve blocklisted functions for the current pull request.
148
+
149
+ Returns A dictionary mapping filenames to sets of blocklisted function names.
150
+ """
151
151
  pr_number = get_pr_number()
152
152
  if pr_number is None:
153
153
  return {}
154
154
 
155
+ not_found = 404
156
+ internal_server_error = 500
157
+
155
158
  owner, repo = get_repo_owner_and_name()
156
- information = {
157
- "pr_number": pr_number,
158
- "repo_owner": owner,
159
- "repo_name": repo,
160
- }
159
+ information = {"pr_number": pr_number, "repo_owner": owner, "repo_name": repo}
161
160
  try:
162
- req = make_cfapi_request(
163
- endpoint="/verify-existing-optimizations",
164
- method="POST",
165
- payload=information,
166
- )
161
+ req = make_cfapi_request(endpoint="/verify-existing-optimizations", method="POST", payload=information)
162
+ if req.status_code == not_found:
163
+ logger.debug(req.json()["message"])
164
+ return {}
165
+ if req.status_code == internal_server_error:
166
+ logger.error(req.json()["message"])
167
+ return {}
168
+ req.raise_for_status()
167
169
  content: dict[str, list[str]] = req.json()
168
170
  except Exception as e:
169
171
  logger.error(f"Error getting blocklisted functions: {e}")
@@ -31,10 +31,7 @@ def parse_args() -> Namespace:
31
31
  init_actions_parser = subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow")
32
32
  init_actions_parser.set_defaults(func=install_github_actions)
33
33
  parser.add_argument("--file", help="Try to optimize only this file")
34
- parser.add_argument(
35
- "--function",
36
- help="Try to optimize only this function within the given file path",
37
- )
34
+ parser.add_argument("--function", help="Try to optimize only this function within the given file path")
38
35
  parser.add_argument(
39
36
  "--all",
40
37
  help="Try to optimize all functions. Can take a really long time. Can pass an optional starting directory to"
@@ -50,30 +47,16 @@ def parse_args() -> Namespace:
50
47
  " This is the top-level root directory where all the Python source code is located.",
51
48
  )
52
49
  parser.add_argument(
53
- "--tests-root",
54
- type=str,
55
- help="Path to the test directory of the project, where all the tests are located.",
50
+ "--tests-root", type=str, help="Path to the test directory of the project, where all the tests are located."
56
51
  )
57
52
  parser.add_argument("--test-framework", choices=["pytest", "unittest"], default="pytest")
53
+ parser.add_argument("--config-file", type=str, help="Path to the pyproject.toml with codeflash configs.")
58
54
  parser.add_argument(
59
- "--config-file",
60
- type=str,
61
- help="Path to the pyproject.toml with codeflash configs.",
55
+ "--use-cached-tests", action="store_true", help="Use cached tests from a specified file for debugging."
62
56
  )
57
+ parser.add_argument("--replay-test", type=str, help="Path to replay test to optimize functions from")
63
58
  parser.add_argument(
64
- "--use-cached-tests",
65
- action="store_true",
66
- help="Use cached tests from a specified file for debugging.",
67
- )
68
- parser.add_argument(
69
- "--replay-test",
70
- type=str,
71
- help="Path to replay test to optimize functions from",
72
- )
73
- parser.add_argument(
74
- "--no-pr",
75
- action="store_true",
76
- help="Do not create a PR for the optimization, only update the code locally.",
59
+ "--no-pr", action="store_true", help="Do not create a PR for the optimization, only update the code locally."
77
60
  )
78
61
  parser.add_argument(
79
62
  "--verify-setup",
@@ -170,6 +153,7 @@ def process_pyproject_config(args: Namespace) -> Namespace:
170
153
  # in this case, the ".." becomes outside project scope, causing issues with un-importable paths
171
154
  args.project_root = project_root_from_module_root(args.module_root, pyproject_file_path)
172
155
  args.tests_root = Path(args.tests_root).resolve()
156
+ args.test_project_root = project_root_from_module_root(args.tests_root, pyproject_file_path)
173
157
  return handle_optimize_all_arg_parsing(args)
174
158
 
175
159
 
@@ -187,7 +171,7 @@ def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
187
171
  except git.exc.InvalidGitRepositoryError:
188
172
  logger.exception(
189
173
  "I couldn't find a git repository in the current directory. "
190
- "I need a git repository to run --all and open PRs for optimizations. Exiting...",
174
+ "I need a git repository to run --all and open PRs for optimizations. Exiting..."
191
175
  )
192
176
  apologize_and_exit()
193
177
  if not args.no_pr and not check_and_push_branch(git_repo):
@@ -10,17 +10,13 @@ import inquirer
10
10
 
11
11
  def apologize_and_exit() -> None:
12
12
  click.echo(
13
- "💡 If you're having trouble, see https://docs.codeflash.ai/getting-started/local-installation for further help getting started with Codeflash!",
13
+ "💡 If you're having trouble, see https://docs.codeflash.ai/getting-started/local-installation for further help getting started with Codeflash!"
14
14
  )
15
15
  click.echo("👋 Exiting...")
16
16
  sys.exit(1)
17
17
 
18
18
 
19
- def inquirer_wrapper(
20
- func: Callable[..., str | bool],
21
- *args: str | bool,
22
- **kwargs: str | bool,
23
- ) -> str | bool:
19
+ def inquirer_wrapper(func: Callable[..., str | bool], *args: str | bool, **kwargs: str | bool) -> str | bool:
24
20
  new_args = []
25
21
  new_kwargs = {}
26
22
 
@@ -29,10 +25,7 @@ def inquirer_wrapper(
29
25
  else:
30
26
  message = kwargs["message"]
31
27
  new_kwargs = kwargs.copy()
32
- split_messages = split_string_to_cli_width(
33
- message,
34
- is_confirm=func == inquirer.confirm,
35
- )
28
+ split_messages = split_string_to_cli_width(message, is_confirm=func == inquirer.confirm)
36
29
  for split_message in split_messages[:-1]:
37
30
  click.echo(split_message)
38
31
 
@@ -77,14 +70,7 @@ def inquirer_wrapper_path(*args: str, **kwargs: str) -> dict[str, str]:
77
70
  new_kwargs["message"] = last_message
78
71
  new_args.append(args[0])
79
72
 
80
- return cast(
81
- dict[str, str],
82
- inquirer.prompt(
83
- [
84
- inquirer.Path(*new_args, **new_kwargs),
85
- ],
86
- ),
87
- )
73
+ return cast(dict[str, str], inquirer.prompt([inquirer.Path(*new_args, **new_kwargs)]))
88
74
 
89
75
 
90
76
  def split_string_to_fit_width(string: str, width: int) -> list[str]:
@@ -21,18 +21,10 @@ from codeflash.api.cfapi import is_github_app_installed_on_repo
21
21
  from codeflash.cli_cmds.cli_common import apologize_and_exit, inquirer_wrapper, inquirer_wrapper_path
22
22
  from codeflash.code_utils.compat import LF
23
23
  from codeflash.code_utils.config_parser import parse_config_file
24
- from codeflash.code_utils.env_utils import (
25
- get_codeflash_api_key,
26
- )
24
+ from codeflash.code_utils.env_utils import get_codeflash_api_key
27
25
  from codeflash.code_utils.git_utils import get_repo_owner_and_name
28
- from codeflash.code_utils.github_utils import (
29
- get_github_secrets_page_url,
30
- require_github_app_or_exit,
31
- )
32
- from codeflash.code_utils.shell_utils import (
33
- get_shell_rc_path,
34
- save_api_key_to_rc,
35
- )
26
+ from codeflash.code_utils.github_utils import get_github_secrets_page_url, require_github_app_or_exit
27
+ from codeflash.code_utils.shell_utils import get_shell_rc_path, save_api_key_to_rc
36
28
  from codeflash.telemetry.posthog_cf import ph
37
29
  from codeflash.version import __version__ as version
38
30
 
@@ -78,12 +70,10 @@ def init_codeflash() -> None:
78
70
  f" codeflash --file <path-to-file> to optimize all functions in a file{LF}"
79
71
  f" codeflash --all to optimize all functions in all files in the module you selected ({setup_info.module_root}){LF}"
80
72
  f"-or-{LF}"
81
- f" codeflash --help to see all options{LF}",
73
+ f" codeflash --help to see all options{LF}"
82
74
  )
83
75
  if did_add_new_key:
84
- click.echo(
85
- "🐚 Don't forget to restart your shell to load the CODEFLASH_API_KEY environment variable!",
86
- )
76
+ click.echo("🐚 Don't forget to restart your shell to load the CODEFLASH_API_KEY environment variable!")
87
77
  click.echo("Or run the following command to reload:")
88
78
  if os.name == "nt":
89
79
  click.echo(f" call {get_shell_rc_path()}")
@@ -111,30 +101,16 @@ def collect_setup_info() -> SetupInfo:
111
101
  curdir = Path.cwd()
112
102
  # Check if the cwd is writable
113
103
  if not os.access(curdir, os.W_OK):
114
- click.echo(
115
- f"❌ The current directory isn't writable, please check your folder permissions and try again.{LF}",
116
- )
104
+ click.echo(f"❌ The current directory isn't writable, please check your folder permissions and try again.{LF}")
117
105
  click.echo("It's likely you don't have write permissions for this folder.")
118
106
  sys.exit(1)
119
107
 
120
108
  # Check for the existence of pyproject.toml or setup.py
121
109
  project_name = check_for_toml_or_setup_file()
122
110
 
123
- ignore_subdirs = [
124
- "venv",
125
- "node_modules",
126
- "dist",
127
- "build",
128
- "build_temp",
129
- "build_scripts",
130
- "env",
131
- "logs",
132
- "tmp",
133
- ]
111
+ ignore_subdirs = ["venv", "node_modules", "dist", "build", "build_temp", "build_scripts", "env", "logs", "tmp"]
134
112
  valid_subdirs = [
135
- d
136
- for d in next(os.walk("."))[1]
137
- if not d.startswith(".") and not d.startswith("__") and d not in ignore_subdirs
113
+ d for d in next(os.walk("."))[1] if not d.startswith(".") and not d.startswith("__") and d not in ignore_subdirs
138
114
  ]
139
115
 
140
116
  valid_module_subdirs = [d for d in valid_subdirs if d != "tests"]
@@ -165,9 +141,7 @@ def collect_setup_info() -> SetupInfo:
165
141
  message="Where are your tests located? "
166
142
  f"(If you don't have any tests yet, I can create an empty tests{os.pathsep} directory for you)",
167
143
  choices=test_subdir_options,
168
- default=(
169
- default_tests_subdir if default_tests_subdir in test_subdir_options else test_subdir_options[0]
170
- ),
144
+ default=(default_tests_subdir if default_tests_subdir in test_subdir_options else test_subdir_options[0]),
171
145
  )
172
146
 
173
147
  if tests_root_answer == create_for_me_option:
@@ -285,11 +259,7 @@ def check_for_toml_or_setup_file() -> str | None:
285
259
  else:
286
260
  if setup_py_path.exists():
287
261
  setup_py_content = setup_py_path.read_text(encoding="utf8")
288
- project_name_match = re.search(
289
- r"setup\s*\([^)]*?name\s*=\s*['\"](.*?)['\"]",
290
- setup_py_content,
291
- re.DOTALL,
292
- )
262
+ project_name_match = re.search(r"setup\s*\([^)]*?name\s*=\s*['\"](.*?)['\"]", setup_py_content, re.DOTALL)
293
263
  if project_name_match:
294
264
  project_name = project_name_match.group(1)
295
265
  click.echo(f"✅ Found setup.py for your project {project_name}")
@@ -300,7 +270,7 @@ def check_for_toml_or_setup_file() -> str | None:
300
270
  click.echo(
301
271
  f"💡 I couldn't find a pyproject.toml in the current directory ({curdir}).{LF}"
302
272
  f"(make sure you're running `codeflash init` from your project's root directory!){LF}"
303
- f"I need this file to store my configuration settings.",
273
+ f"I need this file to store my configuration settings."
304
274
  )
305
275
  ph("cli-no-pyproject-toml-or-setup-py")
306
276
 
@@ -321,14 +291,12 @@ def check_for_toml_or_setup_file() -> str | None:
321
291
 
322
292
  # Check if the pyproject.toml file was created
323
293
  if pyproject_toml_path.exists():
324
- click.echo(
325
- f"✅ Created a pyproject.toml file at {pyproject_toml_path}",
326
- )
294
+ click.echo(f"✅ Created a pyproject.toml file at {pyproject_toml_path}")
327
295
  click.pause()
328
296
  ph("cli-created-pyproject-toml")
329
297
  except OSError:
330
298
  click.echo(
331
- "❌ Failed to create pyproject.toml. Please check your disk permissions and available space.",
299
+ "❌ Failed to create pyproject.toml. Please check your disk permissions and available space."
332
300
  )
333
301
  apologize_and_exit()
334
302
  else:
@@ -341,7 +309,7 @@ def check_for_toml_or_setup_file() -> str | None:
341
309
  def install_github_actions() -> None:
342
310
  try:
343
311
  click.echo(
344
- "⚡️ Codeflash can automatically optimize new Github PRs for you when they're opened. Let's get that set up!",
312
+ "⚡️ Codeflash can automatically optimize new Github PRs for you when they're opened. Let's get that set up!"
345
313
  )
346
314
  config, config_file_path = parse_config_file()
347
315
 
@@ -360,10 +328,7 @@ def install_github_actions() -> None:
360
328
  message=f"I'm going to create a new GitHub actions workflow file at {optimize_yaml_path} ... is this OK?",
361
329
  default=True,
362
330
  )
363
- ph(
364
- "cli-github-optimization-confirm-workflow-creation",
365
- {"confirm_creation": confirm_creation_yes},
366
- )
331
+ ph("cli-github-optimization-confirm-workflow-creation", {"confirm_creation": confirm_creation_yes})
367
332
  if not confirm_creation_yes:
368
333
  click.echo("⏩️ Exiting workflow creation.")
369
334
  ph("cli-github-workflow-skipped")
@@ -374,14 +339,9 @@ def install_github_actions() -> None:
374
339
  py_version = sys.version_info
375
340
  python_version_string = f"'{py_version.major}.{py_version.minor}'"
376
341
  optimize_yml_content = (
377
- files("codeflash")
378
- .joinpath("cli_cmds", "workflows", "codeflash-optimize.yaml")
379
- .read_text(encoding="utf-8")
380
- )
381
- optimize_yml_content = optimize_yml_content.replace(
382
- "{{ python_version }}",
383
- python_version_string,
342
+ files("codeflash").joinpath("cli_cmds", "workflows", "codeflash-optimize.yaml").read_text(encoding="utf-8")
384
343
  )
344
+ optimize_yml_content = optimize_yml_content.replace("{{ python_version }}", python_version_string)
385
345
  with optimize_yaml_path.open("w", encoding="utf8") as optimize_yml_file:
386
346
  optimize_yml_file.write(optimize_yml_content)
387
347
  click.echo(f"✅ Created {optimize_yaml_path}{LF}")
@@ -398,7 +358,7 @@ def install_github_actions() -> None:
398
358
  click.echo(
399
359
  "🐙 I opened your Github secrets page! Note: if you see a 404, you probably don't have access to this "
400
360
  "repo's secrets; ask a repo admin to add it for you, or (not super recommended) you can temporarily "
401
- f"hard-code your api key into the workflow file.{LF}",
361
+ f"hard-code your api key into the workflow file.{LF}"
402
362
  )
403
363
  click.pause()
404
364
  click.echo()
@@ -414,13 +374,13 @@ def install_github_actions() -> None:
414
374
  click.launch(optimize_yaml_path.as_posix())
415
375
  click.echo(
416
376
  "📝 I opened the workflow file in your editor! You'll need to edit the steps that install the right Python "
417
- f"version and any project dependencies. See the comments in the file for more details.{LF}",
377
+ f"version and any project dependencies. See the comments in the file for more details.{LF}"
418
378
  )
419
379
  click.pause()
420
380
  click.echo()
421
381
  click.echo(
422
382
  f"Please commit and push this GitHub actions file to your repo, and you're all set!{LF}"
423
- f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}",
383
+ f"🚀 Codeflash is now configured to automatically optimize new Github PRs!{LF}"
424
384
  )
425
385
  ph("cli-github-workflow-created")
426
386
  except KeyboardInterrupt:
@@ -436,14 +396,12 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
436
396
  except FileNotFoundError:
437
397
  click.echo(
438
398
  f"I couldn't find a pyproject.toml in the current directory.{LF}"
439
- f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file.",
399
+ f"Please create a new empty pyproject.toml file here, OR if you use poetry then run `poetry init`, OR run `codeflash init` again from a directory with an existing pyproject.toml file."
440
400
  )
441
401
  apologize_and_exit()
442
402
 
443
403
  codeflash_section = tomlkit.table()
444
- codeflash_section.add(
445
- tomlkit.comment("All paths are relative to this pyproject.toml's directory."),
446
- )
404
+ codeflash_section.add(tomlkit.comment("All paths are relative to this pyproject.toml's directory."))
447
405
  codeflash_section["module-root"] = setup_info.module_root
448
406
  codeflash_section["tests-root"] = setup_info.tests_root
449
407
  codeflash_section["test-framework"] = setup_info.test_framework
@@ -457,7 +415,7 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
457
415
  elif formatter == "other":
458
416
  formatter_cmds.append("your-formatter $file")
459
417
  click.echo(
460
- "🔧 In pyproject.toml, please replace 'your-formatter' with the command you use to format your code.",
418
+ "🔧 In pyproject.toml, please replace 'your-formatter' with the command you use to format your code."
461
419
  )
462
420
  elif formatter == "don't use a formatter":
463
421
  formatter_cmds.append("disabled")
@@ -483,9 +441,7 @@ def install_github_app() -> None:
483
441
  owner, repo = get_repo_owner_and_name(git_repo)
484
442
 
485
443
  if is_github_app_installed_on_repo(owner, repo):
486
- click.echo(
487
- "🐙 Looks like you've already installed the Codeflash GitHub app on this repository! Continuing…",
488
- )
444
+ click.echo("🐙 Looks like you've already installed the Codeflash GitHub app on this repository! Continuing…")
489
445
 
490
446
  else:
491
447
  click.prompt(
@@ -512,7 +468,7 @@ def install_github_app() -> None:
512
468
  click.echo(
513
469
  f"❌ It looks like the Codeflash GitHub App is not installed on the repository {owner}/{repo}.{LF}"
514
470
  f"You won't be able to create PRs with Codeflash until you install the app.{LF}"
515
- f"In the meantime you can make local only optimizations by using the '--no-pr' flag with codeflash.{LF}",
471
+ f"In the meantime you can make local only optimizations by using the '--no-pr' flag with codeflash.{LF}"
516
472
  )
517
473
  break
518
474
  click.prompt(
@@ -549,15 +505,10 @@ def prompt_api_key() -> bool:
549
505
  existing_api_key = None
550
506
  if existing_api_key:
551
507
  display_key = f"{existing_api_key[:3]}****{existing_api_key[-4:]}"
552
- click.echo(
553
- f"🔑 I found a CODEFLASH_API_KEY in your environment [{display_key}]!",
554
- )
508
+ click.echo(f"🔑 I found a CODEFLASH_API_KEY in your environment [{display_key}]!")
555
509
 
556
510
  use_existing_key = inquirer_wrapper(
557
- inquirer.confirm,
558
- message="Do you want to use this key?",
559
- default=True,
560
- show_default=False,
511
+ inquirer.confirm, message="Do you want to use this key?", default=True, show_default=False
561
512
  )
562
513
  if use_existing_key:
563
514
  ph("cli-existing-api-key-used")
@@ -584,7 +535,7 @@ def enter_api_key_and_save_to_rc() -> None:
584
535
  if not browser_launched:
585
536
  click.echo(
586
537
  f"Opening your Codeflash API key page. Grab a key from there!{LF}"
587
- "You can also open this link manually: https://app.codeflash.ai/app/apikeys",
538
+ "You can also open this link manually: https://app.codeflash.ai/app/apikeys"
588
539
  )
589
540
  click.launch("https://app.codeflash.ai/app/apikeys")
590
541
  browser_launched = True # This does not work on remote consoles
@@ -664,22 +615,11 @@ def test_sort():
664
615
 
665
616
 
666
617
  def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test_path: str) -> None:
667
- command = [
668
- "codeflash",
669
- "--file",
670
- "bubble_sort.py",
671
- "--function",
672
- "sorter",
673
- ]
618
+ command = ["codeflash", "--file", "bubble_sort.py", "--function", "sorter"]
674
619
  sys.stdout.write("Running sample optimization... ")
675
620
  sys.stdout.flush()
676
621
  try:
677
- process = subprocess.run(
678
- command,
679
- text=True,
680
- cwd=args.module_root,
681
- check=False,
682
- )
622
+ process = subprocess.run(command, text=True, cwd=args.module_root, check=False)
683
623
  finally:
684
624
  # Delete the bubble_sort.py file after the test
685
625
  Path(bubble_sort_path).unlink(missing_ok=True)
@@ -688,10 +628,8 @@ def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test
688
628
  click.echo(f"{LF}🗑️ Deleted {bubble_sort_test_path}")
689
629
 
690
630
  if process.returncode == 0:
691
- click.echo(
692
- f"{LF}✅ End-to-end test passed. Codeflash has been correctly set up!",
693
- )
631
+ click.echo(f"{LF}✅ End-to-end test passed. Codeflash has been correctly set up!")
694
632
  else:
695
633
  click.echo(
696
- f"{LF}❌ End-to-end test failed. Please check the logs above, and take a look at https://docs.codeflash.ai/getting-started/local-installation for help and troubleshooting.",
634
+ f"{LF}❌ End-to-end test failed. Please check the logs above, and take a look at https://docs.codeflash.ai/getting-started/local-installation for help and troubleshooting."
697
635
  )