Fix incompatibility with the 'flaky' plugin

pyproject.toml

@@ -39,7 +39,8 @@ optional-dependencies.lint = [
 optional-dependencies.test = [
   "covdefaults>=2.2",
   "pytest>=7.1.2",
-  "coverage>=6.4.4"
+  "coverage>=6.4.4",
+  "flaky>=3.7.0",
 ]
 dynamic = ["version"]
 classifiers = [

src/pytest_memray/plugin.py

@@ -117,14 +117,21 @@ def _build_bin_path() -> Path:
 
         @functools.wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> object | None:
-            result_file = _build_bin_path()
-            with Tracker(result_file):
-                result: object | None = func(*args, **kwargs)
             try:
-                metadata = FileReader(result_file).metadata
-            except OSError:
-                return None
-            self.results[pyfuncitem.nodeid] = Result(metadata, result_file)
+                result_file = _build_bin_path()
+                with Tracker(result_file):
+                    result: object | None = func(*args, **kwargs)
+                try:
+                    metadata = FileReader(result_file).metadata
+                except OSError:
+                    return None
+                self.results[pyfuncitem.nodeid] = Result(metadata, result_file)
+            finally:
+                # Restore the original function. This is needed because some
+                # pytest plugins (e.g. flaky) will call our pytest_pyfunc_call
+                # hook again with whatever is here, which will cause the wrapper
+                # to be wrapped again.
+                pyfuncitem.obj = func
             return result
 
         pyfuncitem.obj = wrapper

tests/test_pytest_memray.py

@@ -181,12 +181,12 @@ def test_bar():
 
     assert "results for test_memray_report.py::test_foo" in output
     assert "Total memory allocated: 2.0KiB" in output
-    assert "valloc:src/memray/_memray_test_utils.pyx" in output
+    assert "valloc:" in output
     assert "-> 2.0KiB" in output
 
     assert "results for test_memray_report.py::test_bar" in output
     assert "Total memory allocated: 1.0KiB" in output
-    assert "valloc:src/memray/_memray_test_utils.pyx" in output
+    assert "valloc:" in output
     assert "-> 1.0KiB" in output
 
 
@@ -223,12 +223,12 @@ def test_bar():
 
     assert "results for test_memray_report.py::test_foo" not in output
     assert "Total memory allocated: 2.0KiB" not in output
-    assert "valloc:src/memray/_memray.pyx" not in output
+    assert "valloc:" not in output
     assert "-> 2.0KiB" not in output
 
     assert "results for test_memray_report.py::test_bar" not in output
     assert "Total memory allocated: 1.0KiB" not in output
-    assert "valloc:src/memray/_memray.pyx" not in output
+    assert "valloc:" not in output
     assert "-> 1.0KiB" not in output
 
 
@@ -368,3 +368,24 @@ def test_t():
 
     assert result.ret == ExitCode.OK
     assert bin_path.exists()
+
+
+def test_plugin_works_with_the_flaky_plugin(pytester: Pytester) -> None:
+    pytester.makepyfile(
+        """
+        from flaky import flaky
+
+        @flaky
+        def test_hello_world():
+            1/0
+        """
+    )
+
+    with patch("pytest_memray.plugin.Tracker") as mock:
+        result = pytester.runpytest("--memray")
+
+    # Ensure that flaky has only called our Tracker once per retry (2 times)
+    # and not more times because it has incorrectly wrapped our plugin and
+    # called it multiple times per retry.
+    assert mock.call_count == 2
+    assert result.ret == ExitCode.TESTS_FAILED