Allow to customize the tracebacks one error (#58)

Co-authored-by: Bernát Gábor <gaborjbernat@gmail.com>

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
 

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.

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__ = [

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,21 +330,41 @@ 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(
         "hide_memray_summary",
         "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")

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",
 ]

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(
         """