diff --git a/README.md b/README.md index 1dc2eba..67b1b72 100644 --- a/README.md +++ b/README.md @@ -94,12 +94,16 @@ MEMORY PROBLEMS demo/test_ok.py::test_memory_exceed temporary folder) - `--memray-bin-prefix` - prefix to use for the binary dump (by default a random UUID4 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) ## Configuration - INI - `memray(bool)` - activate memray tracking - `most-allocations(string)` - show the N tests that allocate most memory (N=0 for all) - `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) ## License diff --git a/docs/news/58.feature.rst b/docs/news/58.feature.rst new file mode 100644 index 0000000..db6e52f --- /dev/null +++ b/docs/news/58.feature.rst @@ -0,0 +1 @@ +Add two new options that allow to customize the ammount of frames in allocation tracebacks as well as including hybrid stack traces. diff --git a/src/pytest_memray/marks.py b/src/pytest_memray/marks.py index 2f2d286..2faa30d 100644 --- a/src/pytest_memray/marks.py +++ b/src/pytest_memray/marks.py @@ -2,11 +2,14 @@ from dataclasses import dataclass from typing import Tuple +from typing import cast from memray import AllocationRecord +from pytest import Config from .utils import parse_memory_string from .utils import sizeof_fmt +from .utils import value_or_ini PytestSection = Tuple[str, str] @@ -18,6 +21,8 @@ class _MemoryInfo: max_memory: float total_allocated_memory: int allocations: list[AllocationRecord] + num_stacks: int + native_stacks: bool @property def section(self) -> PytestSection: @@ -30,11 +35,22 @@ def section(self) -> PytestSection: ] for record in self.allocations: size = record.size - stack_trace = record.stack_trace() + stack_trace = ( + record.hybrid_stack_trace() + if self.native_stacks + else record.stack_trace() + ) if not stack_trace: continue - (function, file, line), *_ = stack_trace - text_lines.append(f"\t- {function}:{file}:{line} -> {sizeof_fmt(size)}") + padding = " " * 4 + text_lines.append(f"{padding}- {sizeof_fmt(size)} allocated here:") + stacks_left = self.num_stacks + for function, file, line in stack_trace: + if stacks_left <= 0: + break + text_lines.append(f"{padding*2}{function}:{file}:{line}") + stacks_left -= 1 + return "memray-max-memory", "\n".join(text_lines) @property @@ -46,14 +62,18 @@ def long_repr(self) -> str: def limit_memory( - limit: str, *, _allocations: list[AllocationRecord] + limit: str, *, _allocations: list[AllocationRecord], _config: Config ) -> _MemoryInfo | None: """Limit memory used by the test.""" max_memory = parse_memory_string(limit) total_allocated_memory = sum(record.size for record in _allocations) if total_allocated_memory < max_memory: return None - return _MemoryInfo(max_memory, total_allocated_memory, _allocations) + num_stacks: int = cast(int, value_or_ini(_config, "stacks")) + native_stacks: bool = cast(bool, value_or_ini(_config, "native")) + return _MemoryInfo( + max_memory, total_allocated_memory, _allocations, num_stacks, native_stacks + ) __all__ = [ diff --git a/src/pytest_memray/plugin.py b/src/pytest_memray/plugin.py index 297c7ce..3b14c83 100644 --- a/src/pytest_memray/plugin.py +++ b/src/pytest_memray/plugin.py @@ -35,7 +35,9 @@ from .marks import limit_memory from .utils import WriteEnabledDirectoryAction +from .utils import positive_int from .utils import sizeof_fmt +from .utils import value_or_ini MARKERS = {"limit_memory": limit_memory} @@ -144,11 +146,13 @@ def _build_bin_path() -> Path: result_file.unlink() return result_file + native: bool = bool(value_or_ini(self.config, "native")) + @functools.wraps(func) def wrapper(*args: Any, **kwargs: Any) -> object | None: try: result_file = _build_bin_path() - with Tracker(result_file): + with Tracker(result_file, native_traces=native): result: object | None = func(*args, **kwargs) try: metadata = FileReader(result_file).metadata @@ -194,8 +198,13 @@ def pytest_runtest_makereport( continue reader = FileReader(result.result_file) func = reader.get_high_watermark_allocation_records - allocations = list(func(merge_threads=True)) - res = marker_fn(*marker.args, **marker.kwargs, _allocations=allocations) + allocations = list((func(merge_threads=True))) + res = marker_fn( + *marker.args, + **marker.kwargs, + _allocations=allocations, + _config=self.config, + ) if res: report.outcome = "failed" report.longrepr = res.long_repr @@ -321,6 +330,19 @@ def pytest_addoption(parser: Parser) -> None: default=5, help="Show the N tests that allocate most memory (N=0 for all)", ) + group.addoption( + "--stacks", + type=positive_int, + default=1, + help="Show the N stack entries when showing tracebacks of memory allocations", + ) + group.addoption( + "--native", + action="store_true", + default=False, + help="Show native frames when showing tracebacks of memory allocations " + "(will be slower)", + ) parser.addini("memray", "Activate pytest.ini setting", type="bool") parser.addini( @@ -328,14 +350,21 @@ def pytest_addoption(parser: Parser) -> None: "Hide the memray summary at the end of the execution", type="bool", ) + parser.addini( + "stacks", + help="Show the N stack entries when showing tracebacks of memory allocations", + type="string", + ) + parser.addini( + "native", + help="Show native frames when showing tracebacks of memory allocations " + "(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) -def value_or_ini(config: Config, key: str) -> object: - return config.getvalue(key) or config.getini(key) - - def pytest_configure(config: Config) -> None: pytest_memray = Manager(config) config.pluginmanager.register(pytest_memray, "memray_manager") diff --git a/src/pytest_memray/utils.py b/src/pytest_memray/utils.py index 907d683..462e66f 100644 --- a/src/pytest_memray/utils.py +++ b/src/pytest_memray/utils.py @@ -1,5 +1,6 @@ from __future__ import annotations +import argparse import os import re from argparse import Action @@ -8,6 +9,8 @@ from pathlib import Path from typing import Sequence +from pytest import Config + def sizeof_fmt(num: int | float, suffix: str = "B") -> str: for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]: @@ -43,6 +46,16 @@ def parse_memory_string(mem_str: str) -> float: return float(quantity) * UNIT_TO_MULTIPLIER[unit.upper()] +def value_or_ini(config: Config, key: str) -> object: + value = config.getvalue(key) + if value: + return value + try: + return config.getini(key) + except (KeyError, ValueError): + return value + + class WriteEnabledDirectoryAction(Action): def __call__( self, @@ -67,8 +80,17 @@ def __call__( setattr(namespace, self.dest, folder) +def positive_int(value: str) -> int: + the_int = int(value) + if the_int <= 0: + raise argparse.ArgumentTypeError(f"{value} is an invalid positive int value") + return the_int + + __all__ = [ "WriteEnabledDirectoryAction", "parse_memory_string", "sizeof_fmt", + "value_or_ini", + "positive_int", ] diff --git a/tests/test_pytest_memray.py b/tests/test_pytest_memray.py index 56a8451..3f75bd6 100644 --- a/tests/test_pytest_memray.py +++ b/tests/test_pytest_memray.py @@ -2,9 +2,11 @@ import xml.etree.ElementTree as ET from types import SimpleNamespace +from unittest.mock import ANY from unittest.mock import patch import pytest +from memray import Tracker from pytest import ExitCode from pytest import Pytester @@ -148,6 +150,67 @@ def test_memory_alloc_fails(): assert result.ret == ExitCode.TESTS_FAILED +@pytest.mark.parametrize("num_stacks", [1, 5, 100]) +def test_memray_report_limit_number_stacks(num_stacks: int, pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + + def rec(n): + if n <= 1: + allocator.valloc(1024*2) + allocator.free() + return None + return rec(n - 1) + + + @pytest.mark.limit_memory("1kb") + def test_foo(): + rec(10) + """ + ) + + result = pytester.runpytest("--memray", f"--stacks={num_stacks}") + + assert result.ret == ExitCode.TESTS_FAILED + + output = result.stdout.str() + + assert "valloc:" in output + assert output.count("rec:") == min(num_stacks - 1, 10) + + +@pytest.mark.parametrize("native", [True, False]) +def test_memray_report_native(native: bool, pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + + @pytest.mark.limit_memory("1kb") + def test_foo(): + allocator.valloc(1024*2) + allocator.free() + """ + ) + + with patch("pytest_memray.plugin.Tracker", wraps=Tracker) as mock: + result = pytester.runpytest("--memray", *(["--native"] if native else [])) + + assert result.ret == ExitCode.TESTS_FAILED + + output = result.stdout.str() + mock.assert_called_once_with(ANY, native_traces=native) + + if native: + assert "MemoryAllocator_1" in output + else: + assert "MemoryAllocator_1" not in output + + def test_memray_report(pytester: Pytester) -> None: pytester.makepyfile( """