Allow specifying the bin path prefix (#31)

And if set recreates it if already exists. Print the number of dumps generated in the summary.

README.md

@@ -92,6 +92,7 @@ MEMORY PROBLEMS demo/test_ok.py::test_memory_exceed
 - `--hide-memray-summary` - hide the memray summary at the end of the execution
 - `--memray-bin-path` - path where to write the memray binary dumps (by default a
   temporary folder)
+- `--memray-bin-prefix` - prefix to use for the binary dump (by default a random UUID4 hex)
 
 ## Configuration - INI
 

docs/configuration.rst

@@ -18,6 +18,9 @@ The complete list of command line options is:
   ``--memray-bin-path``
     Path where to write the memray binary dumps (by default a temporary folder).
 
+  ``--memray-bin-prefix``
+    Prefix to use for the binary dump (by default a random UUID4 hex)
+
 .. tab:: Config file options
 
   ``memray(bool)``

docs/news/28.feature.rst

@@ -0,0 +1,3 @@
+Allow specifying the prefix used for ``-memray-bin-path`` dumps via the
+``-memray-bin-prefix`` (and if specified and file already exists will be recreated) -
+by :user:`gaborbernat`.

src/pytest_memray/plugin.py

@@ -85,30 +85,39 @@ def __init__(self, config: Config) -> None:
         self.results: dict[str, Result] = {}
         self.config = config
         path: Path | None = config.getvalue("memray_bin_path")
-        self.result_path: TemporaryDirectory[str] | Path = path or TemporaryDirectory()
-        self._is_result_tmp: bool = path is None
-        self._run_id = uuid.uuid4().hex
+        if path is None:
+            self._tmp_dir: None | TemporaryDirectory[str] = TemporaryDirectory()
+            self.result_path: Path = Path(self._tmp_dir.name)
+        else:
+            self._tmp_dir = None
+            self.result_path = path
+        self._bin_prefix = config.getvalue("memray_bin_prefix") or uuid.uuid4().hex
 
     @hookimpl(hookwrapper=True)  # type: ignore[misc] # Untyped decorator
     def pytest_unconfigure(self, config: Config) -> Generator[None, None, None]:
         yield
-        if self._is_result_tmp:
-            assert isinstance(self.result_path, TemporaryDirectory)
-            self.result_path.cleanup()
+        if self._tmp_dir is not None:
+            self._tmp_dir.cleanup()
 
     @hookimpl(hookwrapper=True)  # type: ignore[misc] # Untyped decorator
     def pytest_pyfunc_call(self, pyfuncitem: Function) -> object | None:
         func = pyfuncitem.obj
 
-        @functools.wraps(func)
-        def wrapper(*args: Any, **kwargs: Any) -> object | None:
-            if self._is_result_tmp:
-                name = f"{uuid.uuid4().hex}.bin"
-            else:
+        def _build_bin_path() -> Path:
+            if self._tmp_dir is None:
                 of_id = pyfuncitem.nodeid.replace("::", "-")
                 of_id = of_id.replace(os.sep, "-")
-                name = f"{self._run_id}-{of_id}.bin"
-            result_file = Path(self.result_path.name) / name
+                name = f"{self._bin_prefix}-{of_id}.bin"
+            else:
+                name = f"{uuid.uuid4().hex}.bin"
+            result_file = self.result_path / name
+            if self._tmp_dir is None and result_file.exists():
+                result_file.unlink()
+            return result_file
+
+        @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:
@@ -119,7 +128,6 @@ def wrapper(*args: Any, **kwargs: Any) -> object | None:
             return result
 
         pyfuncitem.obj = wrapper
-
         yield
 
     @hookimpl(hookwrapper=True)  # type: ignore[misc] # Untyped decorator
@@ -197,6 +205,10 @@ def pytest_terminal_summary(
                 metadata=reader.metadata,
                 terminalreporter=terminalreporter,
             )
+        if self._tmp_dir is None:
+            msg = f"Created {len(total_sizes)} binary dumps at {self.result_path}"
+            msg += f" with prefix {self._bin_prefix}"
+            terminalreporter.write_line(msg)
 
     @staticmethod
     def _report_records_for_test(
@@ -238,6 +250,11 @@ def pytest_addoption(parser: Parser) -> None:
         default=None,
         help="Path where to write the memray binary dumps (by default a temporary folder)",
     )
+    group.addoption(
+        "--memray-bin-prefix",
+        default=None,
+        help="Prefix to use for the binary dump (by default a random UUID4 hex)",
+    )
     group.addoption(
         "--hide-memray-summary",
         action="store_true",

tests/test_pytest_memray.py

@@ -341,3 +341,30 @@ def test_b(i):
         "H-magic-test_a.py-test_a.bin",
         "H-magic-test_a.py-test_b[1].bin",
     }
+
+    output = result.stdout.str()
+    assert f"Created 3 binary dumps at {dump} with prefix H" in output
+
+
+@pytest.mark.parametrize("override", [True, False])
+def test_bin_path_prefix(pytester: Pytester, override: bool) -> None:
+    py = """
+    import pytest
+    def test_t():
+        assert [1]
+    """
+    pytester.makepyfile(test_a=py)
+
+    bin_path = pytester.path / "p-test_a.py-test_t.bin"
+    if override:
+        bin_path.write_bytes(b"")
+
+    args = ["--memray", "--memray-bin-path", str(pytester.path)]
+    args.extend(["--memray-bin-prefix", "p"])
+    result = pytester.runpytest(*args)
+    res = list(pytester.path.iterdir())
+
+    assert res
+
+    assert result.ret == ExitCode.OK
+    assert bin_path.exists()