From 9fc6c5cc9a0194672a366a49bef82e8ff0186043 Mon Sep 17 00:00:00 2001 From: Pablo Galindo Date: Fri, 23 Jun 2023 14:50:51 +0100 Subject: [PATCH] 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 | 2 ++ docs/configuration.rst | 6 +++++ src/pytest_memray/plugin.py | 20 +++++++++++++++- tests/test_pytest_memray.py | 47 ++++++++++++++++++++++++++++++++++++- 4 files changed, 73 insertions(+), 2 deletions(-) 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..98b9557 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,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) 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( """