Support passing --trace-python-allocators to memray

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,12 @@ The complete list of command line options is:
 
   ``hide_memray_summary(bool)``
     Hide the memray summary at the end of the execution.
+
+  ``stacks(int)``
+    Show the N most recent stack entries when showing tracebacks of memory allocations
+
+  ``native(bool)``
+    Include native frames when showing tracebacks of memory allocations (will be slower)
+
+  ``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(
         """