WIP changes

cli/codeflash/api/aiservice.py

@@ -6,13 +6,13 @@ import platform
 from typing import TYPE_CHECKING, Any
 
 import requests
-from pydantic.dataclasses import dataclass
 from pydantic.json import pydantic_encoder
 
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.env_utils import get_codeflash_api_key
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-from codeflash.models.ExperimentMetadata import ExperimentMetadata
+from codeflash.models.models import OptimizedCandidate
+# from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+# from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.version import __version__ as codeflash_version
 
@@ -23,13 +23,6 @@ if TYPE_CHECKING:
     from codeflash.models.ExperimentMetadata import ExperimentMetadata
 
 
-@dataclass(frozen=True)
-class OptimizedCandidate:
-    source_code: str
-    explanation: str
-    optimization_id: str
-
-
 class AiServiceClient:
     def __init__(self) -> None:
         self.base_url = self.get_aiservice_base_url()

cli/codeflash/code_utils/code_extractor.py

@@ -11,13 +11,13 @@ from libcst.codemod.visitors import AddImportsVisitor, GatherImportsVisitor, Rem
 from libcst.helpers import calculate_module_and_package
 
 from codeflash.cli_cmds.console import logger
-from codeflash.discovery.functions_to_optimize import FunctionParent
+from codeflash.models.models import FunctionSource, FunctionParent
 
 if TYPE_CHECKING:
     from libcst.helpers import ModuleNameAndPackage
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-    from codeflash.models.models import FunctionSource
+
 
 
 class FutureAliasedImportTransformer(cst.CSTTransformer):

cli/codeflash/code_utils/code_replacer.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, TypeVar
 import libcst as cst
 
 from codeflash.code_utils.code_extractor import add_needed_imports_from_module
-from codeflash.discovery.functions_to_optimize import FunctionParent
+from codeflash.models.models import FunctionParent
 
 if TYPE_CHECKING:
     from pathlib import Path

cli/codeflash/code_utils/instrument_existing_tests.py

@@ -8,8 +8,8 @@ import isort
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
-from codeflash.discovery.discover_unit_tests import CodePosition
-from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.models.models import CodePosition, FunctionParent
 
 
 def node_in_call_position(node: ast.stmt, call_positions: list[CodePosition]) -> bool:

cli/codeflash/discovery/discover_unit_tests.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
+import multiprocessing as mp
 import os
 import re
-import sys
 import unittest
 from collections import defaultdict
 from multiprocessing import Process, Queue
@@ -14,33 +14,14 @@ from pydantic.dataclasses import dataclass
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import module_name_from_file_path
+from codeflash.discovery.new_process import run_pytest_discovery_new_process
+from codeflash.models.models import TestsInFile, FunctionCalledInTest, CodePosition
 from codeflash.verification.test_results import TestType
 
 if TYPE_CHECKING:
     from codeflash.verification.verification_utils import TestConfig
 
 
-@dataclass(frozen=True)
-class TestsInFile:
-    test_file: Path
-    test_class: Optional[str]  # This might be unused...
-    test_function: str
-    test_suite: Optional[str]
-    test_type: TestType
-
-
-@dataclass(frozen=True)
-class CodePosition:
-    line_no: int
-    col_no: int
-
-
-@dataclass(frozen=True)
-class FunctionCalledInTest:
-    tests_in_file: TestsInFile
-    position: CodePosition
-
-
 @dataclass(frozen=True)
 class TestFunction:
     function_name: str
@@ -60,51 +41,7 @@ def discover_unit_tests(
     raise ValueError(msg)
 
 
-def run_pytest_discovery_new_process(queue: Queue, cwd: str, tests_root: str) -> tuple[int, list] | None:
-    import pytest
 
-    os.chdir(cwd)
-    collected_tests = []
-    pytest_rootdir: Path | None = None
-    tests: list[TestsInFile] = []
-    sys.path.insert(1, str(cwd))
-
-    class PytestCollectionPlugin:
-        def pytest_collection_finish(self, session) -> None:
-            nonlocal pytest_rootdir
-            collected_tests.extend(session.items)
-            pytest_rootdir = Path(session.config.rootdir)
-
-    try:
-        exitcode = pytest.main(
-            [tests_root, "--collect-only", "-pno:terminal", "-m", "not skip"], plugins=[PytestCollectionPlugin()]
-        )
-    except Exception as e:
-        logger.exception(f"Failed to collect tests: {e!s}")
-        exitcode = -1
-        queue.put((exitcode, tests, pytest_rootdir))
-    tests = parse_pytest_collection_results(collected_tests)
-    queue.put((exitcode, tests, pytest_rootdir))
-
-
-def parse_pytest_collection_results(pytest_tests: str) -> list[TestsInFile]:
-    test_results: list[TestsInFile] = []
-    for test in pytest_tests:
-        test_class = None
-        test_file_path = str(test.path)
-        if test.cls:
-            test_class = test.parent.name
-        test_type = TestType.REPLAY_TEST if "__replay_test" in test_file_path else TestType.EXISTING_UNIT_TEST
-        test_results.append(
-            TestsInFile(
-                test_file=str(test.path),
-                test_class=test_class,
-                test_function=test.name,
-                test_suite=None,  # not used in pytest until now
-                test_type=test_type,
-            )
-        )
-    return test_results
 
 
 def discover_tests_pytest(
@@ -112,6 +49,7 @@ def discover_tests_pytest(
 ) -> dict[str, list[FunctionCalledInTest]]:
     tests_root = cfg.tests_root
     project_root = cfg.project_root_path
+    mp.set_start_method('spawn')
 
     q: Queue = Queue()
     p: Process = Process(target=run_pytest_discovery_new_process, args=(q, project_root, tests_root))
@@ -124,7 +62,7 @@ def discover_tests_pytest(
     else:
         logger.debug(f"Pytest collection exit code: {exitcode}")
     if pytest_rootdir is not None:
-        cfg.tests_project_rootdir = pytest_rootdir
+        cfg.tests_project_rootdir = Path(pytest_rootdir)
     file_to_test_map = defaultdict(list)
     for test in tests:
         if discover_only_these_tests and test.test_file not in discover_only_these_tests:

cli/codeflash/discovery/functions_to_optimize.py

@@ -14,7 +14,7 @@ import libcst as cst
 from pydantic.dataclasses import dataclass
 
 from codeflash.api.cfapi import get_blocklisted_functions
-from codeflash.cli_cmds.console import logger, progress_bar
+from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import (
     is_class_defined_in_file,
     module_name_from_file_path,
@@ -22,6 +22,7 @@ from codeflash.code_utils.code_utils import (
 )
 from codeflash.code_utils.git_utils import get_git_diff
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
+from codeflash.models.models import FunctionParent
 from codeflash.telemetry.posthog_cf import ph
 
 if TYPE_CHECKING:
@@ -102,12 +103,6 @@ class FunctionWithReturnStatement(ast.NodeVisitor):
             self.ast_path.pop()
 
 
-@dataclass(frozen=True)
-class FunctionParent:
-    name: str
-    type: str
-
-
 @dataclass(frozen=True, config={"arbitrary_types_allowed": True})
 class FunctionToOptimize:
     """Represents a function that is a candidate for optimization.

cli/codeflash/discovery/new_process.py

@@ -0,0 +1,54 @@
+import os
+import sys
+from multiprocessing import Queue
+
+#from codeflash.models.models import TestsInFile
+#from codeflash.verification.test_results import TestType
+
+
+def run_pytest_discovery_new_process(queue: Queue, cwd: str, tests_root: str) -> tuple[int, list] | None:
+    sys.modules.pop("returns")
+    import pytest
+
+    os.chdir(cwd)
+    collected_tests = []
+    pytest_rootdir = None
+    tests = []
+    sys.path.insert(1, str(cwd))
+
+    class PytestCollectionPlugin:
+        def pytest_collection_finish(self, session) -> None:
+            nonlocal pytest_rootdir
+            collected_tests.extend(session.items)
+            pytest_rootdir = session.config.rootdir
+
+    try:
+        exitcode = pytest.main(
+            [tests_root, "--collect-only", "-pno:terminal", "-m", "not skip"], plugins=[PytestCollectionPlugin()]
+        )
+    except Exception as e:
+        print(f"Failed to collect tests: {e!s}")
+        exitcode = -1
+        queue.put((exitcode, tests, pytest_rootdir))
+    #tests = parse_pytest_collection_results(collected_tests)
+    queue.put((exitcode, collected_tests, pytest_rootdir))
+
+
+# def parse_pytest_collection_results(pytest_tests: str) -> list[TestsInFile]:
+#     test_results: list[TestsInFile] = []
+#     for test in pytest_tests:
+#         test_class = None
+#         test_file_path = str(test.path)
+#         if test.cls:
+#             test_class = test.parent.name
+#         test_type = TestType.REPLAY_TEST if "__replay_test" in test_file_path else TestType.EXISTING_UNIT_TEST
+#         test_results.append(
+#             TestsInFile(
+#                 test_file=str(test.path),
+#                 test_class=test_class,
+#                 test_function=test.name,
+#                 test_suite=None,  # not used in pytest until now
+#                 test_type=test_type,
+#             )
+#         )
+#     return test_results

cli/codeflash/models/models.py

@@ -7,8 +7,6 @@ from jedi.api.classes import Name
 from pydantic import BaseModel
 from pydantic.dataclasses import dataclass
 
-from codeflash.api.aiservice import OptimizedCandidate
-from codeflash.discovery.functions_to_optimize import FunctionParent
 from codeflash.verification.test_results import TestResults, TestType
 
 
@@ -108,3 +106,37 @@ class OriginalCodeBaseline(BaseModel):
 class OptimizationSet(BaseModel):
     control: list[OptimizedCandidate]
     experiment: Optional[list[OptimizedCandidate]]
+
+
+@dataclass(frozen=True)
+class TestsInFile:
+    test_file: Path
+    test_class: Optional[str]  # This might be unused...
+    test_function: str
+    test_suite: Optional[str]
+    test_type: TestType
+
+
+@dataclass(frozen=True)
+class OptimizedCandidate:
+    source_code: str
+    explanation: str
+    optimization_id: str
+
+
+@dataclass(frozen=True)
+class FunctionCalledInTest:
+    tests_in_file: TestsInFile
+    position: CodePosition
+
+
+@dataclass(frozen=True)
+class CodePosition:
+    line_no: int
+    col_no: int
+
+
+@dataclass(frozen=True)
+class FunctionParent:
+    name: str
+    type: str

cli/codeflash/optimization/function_context.py

@@ -13,8 +13,8 @@ from jedi.api.classes import Name
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_extractor import get_code
 from codeflash.code_utils.code_utils import module_name_from_file_path, path_belongs_to_site_packages
-from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize
-from codeflash.models.models import FunctionSource
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.models.models import FunctionSource, FunctionParent
 
 if TYPE_CHECKING:
     from pathlib import Path

cli/codeflash/optimization/optimizer.py

@@ -39,7 +39,7 @@ from codeflash.code_utils.instrument_existing_tests import inject_profiling_into
 from codeflash.code_utils.remove_generated_tests import remove_functions_from_generated_tests
 from codeflash.code_utils.time_utils import humanize_runtime
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
-from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize, get_functions_to_optimize
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize, get_functions_to_optimize
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import (
     BestOptimization,
@@ -50,7 +50,7 @@ from codeflash.models.models import (
     OptimizedCandidateResult,
     OriginalCodeBaseline,
     TestFile,
-    TestFiles,
+    TestFiles, OptimizedCandidate, FunctionCalledInTest, FunctionParent,
 )
 from codeflash.optimization.function_context import get_constrained_function_context_and_helper_functions
 from codeflash.result.create_pr import check_create_pr, existing_tests_source_for
@@ -69,8 +69,6 @@ if TYPE_CHECKING:
 
     from returns.result import Result
 
-    from codeflash.api.aiservice import OptimizedCandidate
-    from codeflash.discovery.discover_unit_tests import FunctionCalledInTest
     from codeflash.models.models import FunctionSource
 
 

cli/codeflash/result/create_pr.py

@@ -15,8 +15,8 @@ from codeflash.code_utils.git_utils import (
     get_repo_owner_and_name,
     git_root_dir,
 )
-from codeflash.discovery.discover_unit_tests import FunctionCalledInTest
 from codeflash.github.PrComment import FileDiffContent, PrComment
+from codeflash.models.models import FunctionCalledInTest
 from codeflash.result.explanation import Explanation
 
 

cli/tests/test_code_replacement.py

@@ -12,7 +12,8 @@ from codeflash.code_utils.code_replacer import (
     replace_functions_and_add_imports,
     replace_functions_in_file,
 )
-from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.models.models import FunctionParent
 from codeflash.optimization.optimizer import Optimizer
 
 os.environ["CODEFLASH_API_KEY"] = "cf-test-key"

cli/tests/test_function_dependencies.py

@@ -5,7 +5,8 @@ from dataclasses import dataclass
 import pytest
 from returns.pipeline import is_successful
 
-from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.models.models import FunctionParent
 from codeflash.optimization.function_context import get_function_variables_definitions
 from codeflash.optimization.optimizer import Optimizer
 

cli/tests/test_get_code.py

@@ -1,7 +1,8 @@
 import tempfile
 
 from codeflash.code_utils.code_extractor import get_code
-from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.models.models import FunctionParent
 
 
 def test_get_code_function() -> None:

cli/tests/test_get_helper_code.py

@@ -5,7 +5,8 @@ from pathlib import Path
 import pytest
 from returns.pipeline import is_successful
 
-from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.models.models import FunctionParent
 from codeflash.optimization.optimizer import Optimizer
 
 

cli/tests/test_instrumentation.py

@@ -13,9 +13,8 @@ from codeflash.code_utils.instrument_existing_tests import (
     FunctionImportedAsVisitor,
     inject_profiling_into_existing_test,
 )
-from codeflash.discovery.discover_unit_tests import CodePosition
-from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize
-from codeflash.models.models import TestFile, TestFiles
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.models.models import TestFile, TestFiles, CodePosition, FunctionParent
 from codeflash.optimization.optimizer import Optimizer
 from codeflash.verification.test_results import TestType
 

django/aiservice/testgen/gen_inspired_tests.py

@@ -7,12 +7,12 @@ import platform
 from typing import List, Tuple
 
 import openai  # used for calling the OpenAI API
-from codeflash.code_utils.code_extractor import get_code
-from codeflash.code_utils.code_utils import ellipsis_in_ast, get_imports_from_file
-from codeflash.discovery.discover_unit_tests import TestsInFile
 from codeflash.verification.gen_regression_tests import print_message_delta, print_messages
 
 from aiservice.models.aimodels import EXECUTE_MODEL, EXPLAIN_MODEL, LLM, PLAN_MODEL
+from codeflash.code_utils.code_extractor import get_code
+from codeflash.code_utils.code_utils import ellipsis_in_ast, get_imports_from_file
+from codeflash.models.models import TestsInFile
 
 
 def regression_tests_from_function_with_inspiration(