fixup! Add a new marker to check for memory leaks

docs/api.rst

@@ -0,0 +1,15 @@
+.. module:: pytest_memray
+
+pytest-memray API
+=================
+
+Types
+-----
+
+.. autoclass:: StackElement
+
+.. autoclass:: Stack
+   :members:
+
+.. autoclass:: LeaksFilteringFunction
+

docs/conf.py

@@ -9,6 +9,7 @@
 from sphinxcontrib.programoutput import Command
 
 extensions = [
+    "sphinx.ext.autodoc",
     "sphinx.ext.extlinks",
     "sphinx.ext.githubpages",
     "sphinxarg.ext",

docs/index.rst

@@ -14,3 +14,4 @@ reports like:
    usage
    configuration
    news
+   api

docs/usage.rst

@@ -41,6 +41,7 @@ validations on tests when this plugin is enabled.
 ----------------
 
 .. py:function:: limit_memory(memory_limit: str)
+
     Fail the execution of the test if the test allocates more memory than allowed
 
 When this marker is applied to a test, it will cause the test to fail if the execution
@@ -69,12 +70,13 @@ Example of usage:
 ``limit_leaks``
 ---------------
 
- .. py:function:: limit_leaks(location_limit: str, filtering_fn: Callable[Iterable[Tuple[str, str, int]], bool])
-       Fail the execution of the test if any location in the test leaks more memory than allowed.
+.. py:function:: limit_leaks(location_limit: str, filtering_fn: Callable['LeaksFilteringFunction', bool]=None)
+
+    Fail the execution of the test if any call stack in the test leaks more memory than allowed.
 
- .. important::
-       To detect leaks, Memray needs to intercept calls to the Python allocators and use native
-       traces. This is adds significant overhead, and will slow your test down.
+.. important::
+   To detect leaks, Memray needs to intercept calls to the Python allocators and use native
+   traces. This is adds significant overhead, and will slow your test down.
 
 When this marker is applied to a test, it will cause the test to fail if any allocation location in
 the execution of the test leaks more memory than allowed. It takes a single positional argument with a
@@ -90,7 +92,7 @@ The format for the string is ``<NUMBER> ([KMGTP]B|B)``. The marker will raise
 ``ValueError`` if the string format cannot be parsed correctly.
 
 The marker also takes an optional keyword-only argument ``filtering_fn``. This argument represents a filtering
-function that will be called with the traceback for every location that allocates memory that cumulatively is
+function that will be called with the traceback for every call stack that allocates memory that cumulatively is
 bigger than the provided limit. The function must return *True* if the allocation must be taken into account
 and *False* otherwise. This function can be used to discard some false positives detected by the marker.
 
@@ -118,9 +120,9 @@ Example of usage:
 .. warning::
    Is **very** challenging to write tests that do not "leak" memory in some way.
    interpreter caches but there are some that cannot be correctly detected so
-   you may need to allow some small amount of leaked memory per location or use the
+   you may need to allow some small amount of leaked memory per call stack or use the
    ``filtering_fn`` argument to filter out false positive leak reports caused by
    objects that the interpreter plans to reuse later. These caches are
    implementation details of the interpreter, so the amount of memory
-   allocated, the location of the allocation, and the allocator that was used
+   allocated, the call stack of the allocation, and the allocator that was used
    can all change from one Python version to another.

src/pytest_memray/__init__.py

@@ -1,7 +1,13 @@
 from __future__ import annotations
 
 from ._version import __version__ as __version__
+from .marks import StackElement
+from .marks import Stack
+from .marks import LeaksFilteringFunction
 
 __all__ = [
     "__version__",
+    "Stack",
+    "StackElement",
+    "LeaksFilteringFunction",
 ]

src/pytest_memray/marks.py

@@ -5,8 +5,8 @@
 from typing import Tuple
 from typing import cast
 from typing import Callable
-from typing import Iterable
 from typing import Optional
+from typing import Collection
 
 from memray import AllocationRecord
 from memray import FileReader
@@ -17,25 +17,28 @@
 from .utils import value_or_ini
 
 PytestSection = Tuple[str, str]
-
 StackElement = Tuple[str, str, int]
-LeaksFilteringFunction = Callable[[Iterable[StackElement]], bool]
 
 
 @dataclass
-class _MemoryInfo:
+class Stack:
+    frames: Collection[StackElement]
+
+
+LeaksFilteringFunction = Callable[[Stack], bool]
+
+
+@dataclass
+class _MemoryInfoBase:
     """Type that holds memory-related info for a failed test."""
 
     max_memory: float
-    total_allocated_memory: int
     allocations: list[AllocationRecord]
     num_stacks: int
     native_stacks: bool
 
     def _generate_section_text(self, limit_text: str, header_text: str) -> str:
-        text_lines = [
-            f"{header_text} {sizeof_fmt(self.total_allocated_memory)} out of limit of {sizeof_fmt(self.max_memory)}"
-        ]
+        text_lines = [header_text]
         for record in self.allocations:
             size = record.size
             stack_trace = (
@@ -50,18 +53,36 @@ def _generate_section_text(self, limit_text: str, header_text: str) -> str:
             stacks_left = self.num_stacks
             for function, file, line in stack_trace:
                 if stacks_left <= 0:
+                    text_lines.append(f"{padding*2}...")
                     break
                 text_lines.append(f"{padding*2}{function}:{file}:{line}")
                 stacks_left -= 1
 
         return "\n".join(text_lines)
 
+    @property
+    def section(self) -> PytestSection:
+        raise NotImplementedError
+
+    @property
+    def long_repr(self) -> str:
+        raise NotImplementedError
+
+
+@dataclass
+class _MemoryInfo(_MemoryInfoBase):
+    total_allocated_memory: int
+
     @property
     def section(self) -> PytestSection:
         """Return a tuple in the format expected by section reporters."""
+        header_text = (
+            f"List of allocations: {sizeof_fmt(self.total_allocated_memory)} "
+            f"out of limit of {sizeof_fmt(self.max_memory)}"
+        )
         return (
             "memray-max-memory",
-            self._generate_section_text("Test is using", "List of allocations:"),
+            self._generate_section_text("Test is using", header_text),
         )
 
     @property
@@ -71,7 +92,7 @@ def long_repr(self) -> str:
 
 
 @dataclass
-class _LeakedInfo(_MemoryInfo):
+class _LeakedInfo(_MemoryInfoBase):
     """Type that holds leaked memory-related info for a failed test."""
 
     @property
@@ -105,7 +126,7 @@ def limit_memory(
     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
+        max_memory, allocations, num_stacks, native_stacks, total_allocated_memory
     )
 
 
@@ -115,7 +136,7 @@ def limit_leaks(
     filter_fn: Optional[LeaksFilteringFunction] = None,
     _result_file: Path,
     _config: Config,
-) -> _MemoryInfo | None:
+) -> _LeakedInfo | None:
     reader = FileReader(_result_file)
     func = reader.get_leaked_allocation_records
     allocations: list[AllocationRecord] = list((func(merge_threads=True)))
@@ -127,23 +148,19 @@ def limit_leaks(
         for allocation in allocations
         if (
             allocation.size >= memory_limit
-            and (filter_fn is None or filter_fn(allocation.hybrid_stack_trace()))
+            and (filter_fn is None or filter_fn(Stack(allocation.hybrid_stack_trace())))
         )
     )
     if not leaked_allocations:
         return None
-    total_leaked_memory = sum(allocation.size for allocation in leaked_allocations)
-
-    num_stacks: int = cast(int, value_or_ini(_config, "stacks"))
+    sum(allocation.size for allocation in leaked_allocations)
+    num_stacks: int = max(cast(int, value_or_ini(_config, "stacks")), 5)
     return _LeakedInfo(
         memory_limit,
-        total_leaked_memory,
         leaked_allocations,
         num_stacks,
         native_stacks=True,
     )
 
 
-__all__ = [
-    "limit_memory",
-]
+__all__ = ["limit_memory", "limit_leaks", "Stack"]

src/pytest_memray/plugin.py

@@ -152,9 +152,7 @@ def pytest_pyfunc_call(self, pyfuncitem: Function) -> object | None:
             return
 
         if len(markers) > 1:
-            raise ValueError(
-                "Only one memray marker can be applied at the same time to the same test"
-            )
+            raise ValueError("Only one Memray marker can be applied to each test")
 
         def _build_bin_path() -> Path:
             if self._tmp_dir is None and not os.getenv("MEMRAY_RESULT_PATH"):

tests/test_pytest_memray.py

@@ -670,8 +670,8 @@ def this_should_not_be_there():
              allocator.valloc(LEAK_SIZE)
              # No free call here
 
-         def filtering_function(locations):
-             for fn, _, _ in locations:
+         def filtering_function(stack):
+             for fn, _, _ in stack.frames:
                  if fn == "this_should_not_be_there":
                      return False
              return True
@@ -721,7 +721,7 @@ def test_bar():
     assert result.ret == ExitCode.TESTS_FAILED
 
     output = result.stdout.str()
-    assert "Only one memray marker can be applied" in output
+    assert "Only one Memray marker can be applied to each test" in output
 
 
 def test_multiple_markers_are_not_supported_with_global_marker(
@@ -741,4 +741,4 @@ def test_bar():
     assert result.ret == ExitCode.TESTS_FAILED
 
     output = result.stdout.str()
-    assert "Only one memray marker can be applied" in output
+    assert "Only one Memray marker can be applied to each test" in output