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

@@ -40,6 +40,10 @@ validations on tests when this plugin is enabled.
 ``limit_memory``
 ----------------
 
+.. 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
 of the test allocates more memory than allowed. It takes a single argument with a
 string indicating the maximum memory that the test can allocate.
@@ -61,3 +65,64 @@ Example of usage:
     @pytest.mark.limit_memory("24 MB")
     def test_foobar():
         pass  # do some stuff that allocates memory
+
+
+``limit_leaks``
+---------------
+
+.. 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.
+
+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
+string indicating the maximum memory **per allocation location** that the test is allowed to leak.
+
+Leaks are defined as memory that is allocated **in the marked test** that is not freed before leaving the test body.
+
+.. important::
+    It's recommended to run your API or code in a loop when utilizing this plugin. This practice helps in distinguishing
+    genuine leaks from the "noise" generated by internal caches and other incidental allocations.
+
+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 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.
+
+.. tip::
+
+   You can pass the ``--memray-bin-path`` argument to ``pytest`` to specify
+   a directory where Memray will store the binary files with the results. You
+   can then use the ``memray`` CLI to further investigate the allocations and the
+   leaks using any Memray reporters you'd like. Check `the memray docs
+   <https://bloomberg.github.io/memray/getting_started.html>`_ for more
+   information.
+
+Example of usage:
+
+.. code-block:: python
+
+    @pytest.mark.limit_leaks("1 MB")
+    def test_foobar():
+        # Run the function to test in a loop to ensure
+        # we can differentiate leaks from memory allocated
+        # in internal caches
+        for _ in range(100):
+            do_some_stuff()
+
+.. 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 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 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

@@ -1,38 +1,44 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Tuple
 from typing import cast
+from typing import Callable
+from typing import Optional
+from typing import Collection
 
 from memray import AllocationRecord
+from memray import FileReader
 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]
+StackElement = Tuple[str, str, int]
 
 
 @dataclass
-class _MemoryInfo:
-    """Type that holds all memray-related info for a failed test."""
+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
 
-    @property
-    def section(self) -> PytestSection:
-        """Return a tuple in the format expected by section reporters."""
-        total_memory_str = sizeof_fmt(self.total_allocated_memory)
-        max_memory_str = sizeof_fmt(self.max_memory)
-        text_lines = [
-            f"Test is using {total_memory_str} out of limit of {max_memory_str}",
-            "List of allocations: ",
-        ]
+    def _generate_section_text(self, limit_text: str, header_text: str) -> str:
+        text_lines = [header_text]
         for record in self.allocations:
             size = record.size
             stack_trace = (
@@ -47,35 +53,114 @@ def section(self) -> PytestSection:
             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 "memray-max-memory", "\n".join(text_lines)
+        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", header_text),
+        )
+
+    @property
+    def long_repr(self) -> str:
+        """Generate a longrepr user-facing error message."""
+        return f"Test was limited to {sizeof_fmt(self.max_memory)} but allocated {sizeof_fmt(self.total_allocated_memory)}"
+
+
+@dataclass
+class _LeakedInfo(_MemoryInfoBase):
+    """Type that holds leaked memory-related info for a failed test."""
+
+    @property
+    def section(self) -> PytestSection:
+        """Return a tuple in the format expected by section reporters."""
+        return (
+            "memray-leaked-memory",
+            self._generate_section_text("Test leaked", "List of leaked allocations:"),
+        )
 
     @property
     def long_repr(self) -> str:
         """Generate a longrepr user-facing error message."""
-        total_memory_str = sizeof_fmt(self.total_allocated_memory)
-        max_memory_str = sizeof_fmt(self.max_memory)
-        return f"Test was limited to {max_memory_str} but allocated {total_memory_str}"
+        return (
+            f"Test was allowed to leak {sizeof_fmt(self.max_memory)} "
+            "per location but at least one location leaked more"
+        )
 
 
 def limit_memory(
-    limit: str, *, _allocations: list[AllocationRecord], _config: Config
+    limit: str, *, _result_file: Path, _config: Config
 ) -> _MemoryInfo | None:
     """Limit memory used by the test."""
+    reader = FileReader(_result_file)
+    func = reader.get_high_watermark_allocation_records
+    allocations: list[AllocationRecord] = list((func(merge_threads=True)))
     max_memory = parse_memory_string(limit)
-    total_allocated_memory = sum(record.size for record in _allocations)
+    total_allocated_memory = sum(record.size for record in allocations)
     if total_allocated_memory < max_memory:
         return None
     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
+    )
+
+
+def limit_leaks(
+    location_limit: str,
+    *,
+    filter_fn: Optional[LeaksFilteringFunction] = None,
+    _result_file: Path,
+    _config: Config,
+) -> _LeakedInfo | None:
+    reader = FileReader(_result_file)
+    func = reader.get_leaked_allocation_records
+    allocations: list[AllocationRecord] = list((func(merge_threads=True)))
+
+    memory_limit = parse_memory_string(location_limit)
+
+    leaked_allocations = list(
+        allocation
+        for allocation in allocations
+        if (
+            allocation.size >= memory_limit
+            and (filter_fn is None or filter_fn(Stack(allocation.hybrid_stack_trace())))
+        )
+    )
+    if not leaked_allocations:
+        return None
+    sum(allocation.size for allocation in leaked_allocations)
+    num_stacks: int = max(cast(int, value_or_ini(_config, "stacks")), 5)
+    return _LeakedInfo(
+        memory_limit,
+        leaked_allocations,
+        num_stacks,
+        native_stacks=True,
     )
 
 
-__all__ = [
-    "limit_memory",
-]
+__all__ = ["limit_memory", "limit_leaks", "Stack"]

src/pytest_memray/plugin.py

@@ -17,6 +17,7 @@
 from typing import List
 from typing import Tuple
 from typing import cast
+from typing import Protocol
 
 from _pytest.terminal import TerminalReporter
 from memray import AllocationRecord
@@ -34,12 +35,28 @@
 from pytest import hookimpl
 
 from .marks import limit_memory
+from .marks import _MemoryInfo
+from .marks import limit_leaks
 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}
+
+class PluginFn(Protocol):
+    def __call__(
+        *args: Any,
+        _result_file: Path,
+        _config: Config,
+        **kwargs: Any,
+    ) -> _MemoryInfo | None:
+        ...
+
+
+MARKERS = {
+    "limit_memory": limit_memory,
+    "limit_leaks": limit_leaks,
+}
 
 N_TOP_ALLOCS = 5
 N_HISTOGRAM_BINS = 5
@@ -134,6 +151,9 @@ def pytest_pyfunc_call(self, pyfuncitem: Function) -> object | None:
             yield
             return
 
+        if len(markers) > 1:
+            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"):
                 of_id = pyfuncitem.nodeid.replace("::", "-")
@@ -151,6 +171,9 @@ def _build_bin_path() -> Path:
             value_or_ini(self.config, "trace_python_allocators")
         )
 
+        if markers and "limit_leaks" in markers:
+            native = trace_python_allocators = True
+
         @functools.wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> object | None:
             test_result: object | Any = None
@@ -198,19 +221,17 @@ def pytest_runtest_makereport(
             return None
 
         for marker in item.iter_markers():
-            marker_fn = MARKERS.get(marker.name)
-            if not marker_fn:
+            maybe_marker_fn = MARKERS.get(marker.name)
+            if not maybe_marker_fn:
                 continue
+            marker_fn: PluginFn = cast(PluginFn, maybe_marker_fn)
             result = self.results.get(item.nodeid)
             if not result:
                 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,
+                _result_file=result.result_file,
                 _config=self.config,
             )
             if res: