Support passing --trace-python-allocators to memray

Several users have indicated that passing the `--trace-python-allocators`
flag could be very useful for they use cases as they don't want to be
affected by the global state of pymalloc or because they want to
include all the requests for memory at the Python layer so the plugin
behaves more similarly as tracemalloc.

README.md

@@ -96,6 +96,7 @@ MEMORY PROBLEMS demo/test_ok.py::test_memory_exceed
   hex)
 --stacks=STACKS - Show the N stack entries when showing tracebacks of memory allocations
 --native - Show native frames when showing tracebacks of memory allocations (will be slower)
+--trace-python-allocators - Record allocations made by the Pymalloc allocator (will be slower)
 
 ## Configuration - INI
 
@@ -104,6 +105,7 @@ MEMORY PROBLEMS demo/test_ok.py::test_memory_exceed
 - `hide_memray_summary(bool)` - hide the memray summary at the end of the execution
 - `stacks(int)` - Show the N stack entries when showing tracebacks of memory allocations
 - `native(bool)`- Show native frames when showing tracebacks of memory allocations (will be slower)
+- `trace_python_allocators` - Record allocations made by the Pymalloc allocator (will be slower)
 
 ## License
 

docs/configuration.rst

@@ -27,6 +27,9 @@ The complete list of command line options is:
   ``--native``
     Include native frames when showing tracebacks of memory allocations (will be slower)
 
+  ``--trace-python-allocators``
+    Record allocations made by the Pymalloc allocator (will be slower)
+
 .. tab:: Config file options
 
   ``memray(bool)``
@@ -37,3 +40,6 @@ The complete list of command line options is:
 
   ``hide_memray_summary(bool)``
     Hide the memray summary at the end of the execution.
+
+  ``trace_python_allocators(bool)``
+    Record allocations made by the Pymalloc allocator (will be slower)

src/pytest_memray/plugin.py

@@ -147,13 +147,20 @@ def _build_bin_path() -> Path:
             return result_file
 
         native: bool = bool(value_or_ini(self.config, "native"))
+        trace_python_allocators: bool = bool(
+            value_or_ini(self.config, "trace_python_allocators")
+        )
 
         @functools.wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> object | None:
             test_result: object | Any = None
             try:
                 result_file = _build_bin_path()
-                with Tracker(result_file, native_traces=native):
+                with Tracker(
+                    result_file,
+                    native_traces=native,
+                    trace_python_allocators=trace_python_allocators,
+                ):
                     test_result = func(*args, **kwargs)
                 try:
                     metadata = FileReader(result_file).metadata
@@ -344,6 +351,12 @@ def pytest_addoption(parser: Parser) -> None:
         help="Show native frames when showing tracebacks of memory allocations "
         "(will be slower)",
     )
+    group.addoption(
+        "--trace-python-allocators",
+        action="store_true",
+        default=False,
+        help="Record allocations made by the Pymalloc allocator (will be slower)",
+    )
 
     parser.addini("memray", "Activate pytest.ini setting", type="bool")
     parser.addini(
@@ -362,6 +375,11 @@ def pytest_addoption(parser: Parser) -> None:
         "(will be slower)",
         type="bool",
     )
+    parser.addini(
+        "trace_python_allocators",
+        help="Record allocations made by the Pymalloc allocator (will be slower)",
+        type="bool",
+    )
     help_msg = "Show the N tests that allocate most memory (N=0 for all)"
     parser.addini("most_allocations", help_msg)
 

tests/test_pytest_memray.py

@@ -203,14 +203,59 @@ def test_foo():
     assert result.ret == ExitCode.TESTS_FAILED
 
     output = result.stdout.str()
-    mock.assert_called_once_with(ANY, native_traces=native)
+    mock.assert_called_once_with(
+        ANY, native_traces=native, trace_python_allocators=False
+    )
 
     if native:
         assert "MemoryAllocator_1" in output
     else:
         assert "MemoryAllocator_1" not in output
 
 
+@pytest.mark.parametrize("trace_python_allocators", [True, False])
+def test_memray_report_python_allocators(
+    trace_python_allocators: bool, pytester: Pytester
+) -> None:
+    pytester.makepyfile(
+        """
+        import pytest
+        from memray._test import PymallocMemoryAllocator
+        from memray._test import PymallocDomain
+
+        allocator = PymallocMemoryAllocator(PymallocDomain.PYMALLOC_OBJECT)
+
+        def allocate_with_pymalloc():
+            allocator.malloc(256)
+            allocator.free()
+
+        @pytest.mark.limit_memory("128B")
+        def test_foo():
+            allocate_with_pymalloc()
+    """
+    )
+
+    with patch("pytest_memray.plugin.Tracker", wraps=Tracker) as mock:
+        result = pytester.runpytest(
+            "--memray",
+            *(["--trace-python-allocators"] if trace_python_allocators else []),
+        )
+
+    assert result.ret == (
+        ExitCode.TESTS_FAILED if trace_python_allocators else ExitCode.OK
+    )
+
+    output = result.stdout.str()
+    mock.assert_called_once_with(
+        ANY, native_traces=False, trace_python_allocators=trace_python_allocators
+    )
+
+    if trace_python_allocators:
+        assert "allocate_with_pymalloc" in output
+    else:
+        assert "allocate_with_pymalloc" not in output
+
+
 def test_memray_report(pytester: Pytester) -> None:
     pytester.makepyfile(
         """