diff --git a/README.md b/README.md index 67b1b72..aaf493e 100644 --- a/README.md +++ b/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 diff --git a/docs/configuration.rst b/docs/configuration.rst index 1081c74..18dc789 100644 --- a/docs/configuration.rst +++ b/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) diff --git a/src/pytest_memray/plugin.py b/src/pytest_memray/plugin.py index 793c2c1..1933388 100644 --- a/src/pytest_memray/plugin.py +++ b/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) diff --git a/tests/test_pytest_memray.py b/tests/test_pytest_memray.py index e1b7f70..dd19afb 100644 --- a/tests/test_pytest_memray.py +++ b/tests/test_pytest_memray.py @@ -203,7 +203,9 @@ 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 @@ -211,6 +213,49 @@ def test_foo(): 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( """