From 7705a5349ba2144ecd55d90ccb820476da3cfdce Mon Sep 17 00:00:00 2001
From: Matt Wozniski <mwozniski@bloomberg.net>
Date: Thu, 17 Aug 2023 20:03:26 -0400
Subject: [PATCH] fixup! fixup! fixup! Add a new marker to check for memory
 leaks

---
 docs/api.rst                  |  5 +++--
 src/pytest_memray/__init__.py |  4 ++--
 src/pytest_memray/marks.py    | 41 +++++++++++++++++++++++++++++++----
 tests/test_pytest_memray.py   |  4 ++--
 4 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/docs/api.rst b/docs/api.rst
index e8a230d..c40e1c6 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -6,9 +6,10 @@ pytest-memray API
 Types
 -----
 
-.. autoclass:: StackElement
+.. autoclass:: StackFrame()
+   :members:
 
-.. autoclass:: Stack
+.. autoclass:: Stack()
    :members:
 
 .. autoclass:: LeaksFilteringFunction
diff --git a/src/pytest_memray/__init__.py b/src/pytest_memray/__init__.py
index 021f836..02670d4 100644
--- a/src/pytest_memray/__init__.py
+++ b/src/pytest_memray/__init__.py
@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 from ._version import __version__ as __version__
-from .marks import StackElement
+from .marks import StackFrame
 from .marks import Stack
 from .marks import LeaksFilteringFunction
 
 __all__ = [
     "__version__",
     "Stack",
-    "StackElement",
+    "StackFrame",
     "LeaksFilteringFunction",
 ]
diff --git a/src/pytest_memray/marks.py b/src/pytest_memray/marks.py
index 14edc00..c51a191 100644
--- a/src/pytest_memray/marks.py
+++ b/src/pytest_memray/marks.py
@@ -2,11 +2,12 @@
 
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Iterable
 from typing import Tuple
 from typing import cast
 from typing import Callable
 from typing import Optional
-from typing import Collection
+from typing import Sequence
 
 from memray import AllocationRecord
 from memray import FileReader
@@ -17,12 +18,34 @@
 from .utils import value_or_ini
 
 PytestSection = Tuple[str, str]
-StackElement = Tuple[str, str, int]
+
+
+@dataclass
+class StackFrame:
+    """One frame of a call stack.
+
+    Each frame has attributes to tell you what code was executing.
+    """
+
+    function: str
+    """The function being executed, or ``"???"`` if unknown."""
+
+    filename: str
+    """The source file being executed, or ``"???"`` if unknown."""
+
+    lineno: int
+    """The line number of the executing line, or ``0`` if unknown."""
 
 
 @dataclass
 class Stack:
-    frames: Collection[StackElement]
+    """The call stack that led to some memory allocation.
+
+    You can inspect the frames which make up the call stack.
+    """
+
+    frames: Sequence[StackFrame]
+    """The frames that make up the call stack, most recent first."""
 
 
 LeaksFilteringFunction = Callable[[Stack], bool]
@@ -112,6 +135,16 @@ def long_repr(self) -> str:
         )
 
 
+def passes_filter(
+    stack: Iterable[Tuple[str, str, int]], filter_fn: Optional[LeaksFilteringFunction]
+) -> bool:
+    if filter_fn is None:
+        return True
+
+    frames = [StackFrame(*frame) for frame in stack]
+    return filter_fn(Stack(frames))
+
+
 def limit_memory(
     limit: str, *, _result_file: Path, _config: Config
 ) -> _MemoryInfo | None:
@@ -148,7 +181,7 @@ def limit_leaks(
         for allocation in allocations
         if (
             allocation.size >= memory_limit
-            and (filter_fn is None or filter_fn(Stack(allocation.hybrid_stack_trace())))
+            and passes_filter(allocation.hybrid_stack_trace(), filter_fn)
         )
     )
     if not leaked_allocations:
diff --git a/tests/test_pytest_memray.py b/tests/test_pytest_memray.py
index ae9e0e8..f95ce76 100644
--- a/tests/test_pytest_memray.py
+++ b/tests/test_pytest_memray.py
@@ -671,8 +671,8 @@ def this_should_not_be_there():
              # No free call here
 
          def filtering_function(stack):
-             for fn, _, _ in stack.frames:
-                 if fn == "this_should_not_be_there":
+             for frame in stack.frames:
+                 if frame.function == "this_should_not_be_there":
                      return False
              return True