better impl

Author: KRRT7

codeflash/code_utils/codeflash_wrap_decorator.py

@@ -1,23 +1,17 @@
 from __future__ import annotations
 
 import asyncio
-import contextlib
 import gc
-import inspect
 import os
 import sqlite3
-import time
 from enum import Enum
 from functools import wraps
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any, Callable, TypeVar
+from typing import Any, Callable, TypeVar
 
 import dill as pickle
 
-if TYPE_CHECKING:
-    from types import FrameType
-
 
 class VerificationType(str, Enum):  # moved from codeflash/verification/codeflash_capture.py
     FUNCTION_CALL = (
@@ -36,175 +30,19 @@ def get_run_tmp_file(file_path: Path) -> Path:  # moved from codeflash/code_util
     return Path(get_run_tmp_file.tmpdir.name) / file_path
 
 
-def _extract_class_name_tracer(frame_locals: dict[str, Any]) -> str | None:
-    try:
-        self_arg = frame_locals.get("self")
-        if self_arg is not None:
-            try:
-                return self_arg.__class__.__name__
-            except (AttributeError, Exception):
-                cls_arg = frame_locals.get("cls")
-                if cls_arg is not None:
-                    with contextlib.suppress(AttributeError, Exception):
-                        return cls_arg.__name__
-        else:
-            cls_arg = frame_locals.get("cls")
-            if cls_arg is not None:
-                with contextlib.suppress(AttributeError, Exception):
-                    return cls_arg.__name__
-    except Exception:
-        return None
-    return None
-
-
-def _get_module_name_cf_tracer(frame: FrameType | None) -> str:
-    try:
-        test_module = inspect.getmodule(frame)
-    except Exception:
-        test_module = None
-
-    if test_module is not None:
-        module_name = getattr(test_module, "__name__", None)
-        if module_name is not None:
-            return module_name
-
-    if frame is not None:
-        return frame.f_globals.get("__name__", "unknown_module")
-    return "unknown_module"
+def extract_test_context_from_frame() -> tuple[str, str | None, str]:
+    # test_module = os.environ.get("CODEFLASH_TEST_MODULE")
+    test_module = os.environ["CODEFLASH_TEST_MODULE"]
+    test_class = os.environ.get("CODEFLASH_TEST_CLASS", None)
+    # test_function = os.environ.get("CODEFLASH_TEST_FUNCTION")
+    test_function = os.environ["CODEFLASH_TEST_FUNCTION"]
 
+    if test_module and test_function:
+        return (test_module, test_class if test_class else None, test_function)
 
-def extract_test_context_from_frame() -> tuple[str, str | None, str]:
-    frame = inspect.currentframe()
-    # optimize?
-    try:
-        frames_info = []
-        potential_tests = []
-
-        # First pass: collect all frame information
-        if frame is not None:
-            frame = frame.f_back
-
-        while frame is not None:
-            try:
-                function_name = frame.f_code.co_name
-                filename = frame.f_code.co_filename
-                filename_path = Path(filename)
-                frame_locals = frame.f_locals
-                test_module_name = _get_module_name_cf_tracer(frame)
-                class_name = _extract_class_name_tracer(frame_locals)
-
-                frames_info.append(
-                    {
-                        "function_name": function_name,
-                        "filename_path": filename_path,
-                        "frame_locals": frame_locals,
-                        "test_module_name": test_module_name,
-                        "class_name": class_name,
-                        "frame": frame,
-                    }
-                )
-
-            except Exception:  # noqa: S112
-                continue
-
-            frame = frame.f_back
-
-        # Second pass: analyze frames with full context
-        test_class_candidates = []
-        for frame_info in frames_info:
-            function_name = frame_info["function_name"]
-            filename_path = frame_info["filename_path"]
-            frame_locals = frame_info["frame_locals"]
-            test_module_name = frame_info["test_module_name"]
-            class_name = frame_info["class_name"]
-            frame_obj = frame_info["frame"]
-
-            # Keep track of test classes
-            if class_name and (
-                class_name.startswith("Test") or class_name.endswith("Test") or "test" in class_name.lower()
-            ):
-                test_class_candidates.append((class_name, test_module_name))
-
-        # Now process frames again looking for test functions with full candidates list
-        # Collect all test functions to prioritize outer ones over nested ones
-        test_functions = []
-        for frame_info in frames_info:
-            function_name = frame_info["function_name"]
-            filename_path = frame_info["filename_path"]
-            frame_locals = frame_info["frame_locals"]
-            test_module_name = frame_info["test_module_name"]
-            class_name = frame_info["class_name"]
-            frame_obj = frame_info["frame"]
-
-            # Collect test functions
-            if function_name.startswith("test_"):
-                test_class_name = class_name
-
-                # If no class found in current frame, check if we have any test class candidates
-                # Prefer the innermost (first) test class candidate which is more specific
-                if test_class_name is None and test_class_candidates:
-                    test_class_name = test_class_candidates[0][0]
-
-                test_functions.append((test_module_name, test_class_name, function_name))
-
-        # Prioritize test functions with class context, then innermost
-        if test_functions:
-            # First prefer test functions with class context
-            for test_func in test_functions:
-                if test_func[1] is not None:  # has class_name
-                    return test_func
-            # If no test function has class context, return the outermost (most likely the actual test method)
-            return test_functions[-1]
-
-        # If no direct test functions found, look for other test patterns
-        for frame_info in frames_info:
-            function_name = frame_info["function_name"]
-            filename_path = frame_info["filename_path"]
-            frame_locals = frame_info["frame_locals"]
-            test_module_name = frame_info["test_module_name"]
-            class_name = frame_info["class_name"]
-            frame_obj = frame_info["frame"]
-
-            # Test file/module detection
-            if (
-                frame_obj.f_globals.get("__name__", "").startswith("test_")
-                or filename_path.stem.startswith("test_")
-                or "test" in filename_path.parts
-            ):
-                if class_name and (
-                    class_name.startswith("Test") or class_name.endswith("Test") or "test" in class_name.lower()
-                ):
-                    potential_tests.append((test_module_name, class_name, function_name))
-                elif "test" in test_module_name or filename_path.stem.startswith("test_"):
-                    # For functions without class context, try to find the most recent test class
-                    best_class = test_class_candidates[0][0] if test_class_candidates else None
-                    potential_tests.append((test_module_name, best_class, function_name))
-
-            # Framework integration detection
-            if (
-                (
-                    function_name in ["runTest", "_runTest", "run", "_testMethodName"]
-                    or "pytest" in str(frame_obj.f_globals.get("__file__", ""))
-                    or "unittest" in str(frame_obj.f_globals.get("__file__", ""))
-                )
-                and class_name
-                and (class_name.startswith("Test") or "test" in class_name.lower())
-            ):
-                test_method = function_name
-                if "self" in frame_locals:
-                    with contextlib.suppress(AttributeError, TypeError):
-                        test_method = getattr(frame_locals["self"], "_testMethodName", function_name)
-                potential_tests.append((test_module_name, class_name, test_method))
-
-        if potential_tests:
-            for test_module, test_class, test_func in potential_tests:
-                if test_func.startswith("test_"):
-                    return test_module, test_class, test_func
-            return potential_tests[0]
-
-        raise RuntimeError("No test function found in call stack")
-    finally:
-        del frame
+    raise RuntimeError(
+        "Test context environment variables not set - ensure tests are run through codeflash test runner"
+    )
 
 
 def codeflash_behavior_async(func: F) -> F:

codeflash/verification/pytest_plugin.py

@@ -450,3 +450,25 @@ def make_progress_id(i: int, n: int = count) -> str:
             metafunc.parametrize(
                 "__pytest_loop_step_number", range(count), indirect=True, ids=make_progress_id, scope=scope
             )
+
+    @pytest.hookimpl(tryfirst=True)
+    def pytest_runtest_setup(self, item: pytest.Item) -> None:
+        test_module_name = item.module.__name__ if item.module else "unknown_module"
+
+        test_class_name = None
+        if item.cls:
+            test_class_name = item.cls.__name__
+
+        test_function_name = item.name
+        if "[" in test_function_name:
+            test_function_name = test_function_name.split("[", 1)[0]
+
+        os.environ["CODEFLASH_TEST_MODULE"] = test_module_name
+        os.environ["CODEFLASH_TEST_CLASS"] = test_class_name or ""
+        os.environ["CODEFLASH_TEST_FUNCTION"] = test_function_name
+
+    @pytest.hookimpl(trylast=True)
+    def pytest_runtest_teardown(self, _: pytest.Item) -> None:
+        """Clean up test context environment variables after each test."""
+        for var in ["CODEFLASH_TEST_MODULE", "CODEFLASH_TEST_CLASS", "CODEFLASH_TEST_FUNCTION"]:
+            os.environ.pop(var, None)