Add a new current_thread_only to all markers

docs/news/117.feature.rst

@@ -0,0 +1,3 @@
+Add a new ``current_thread_only`` keyword argument to the ``limit_memory`` and
+``limit_leaks`` markers to ignore all allocations made in threads other than
+the one running the test.

docs/usage.rst

@@ -35,7 +35,7 @@ This plugin provides `markers <https://docs.pytest.org/en/latest/example/markers
 that can be used to enforce additional checks and validations on tests.
 
 
-.. py:function:: pytest.mark.limit_memory(memory_limit: str)
+.. py:function:: pytest.mark.limit_memory(memory_limit: str, current_thread_only: bool = False)
 
     Fail the execution of the test if the test allocates more peak memory than allowed.
 
@@ -47,6 +47,10 @@ that can be used to enforce additional checks and validations on tests.
     The format for the string is ``<NUMBER> ([KMGTP]B|B)``. The marker will raise
     ``ValueError`` if the string format cannot be parsed correctly.
 
+    If the optional keyword-only argument ``current_thread_only`` is set to *True*, the
+    plugin will only track memory allocations made by the current thread and all other
+    allocations will be ignored.
+
     .. warning::
 
         As the Python interpreter has its own
@@ -96,6 +100,10 @@ that can be used to enforce additional checks and validations on tests.
     ignored, the test will not fail. This can be used to discard any known false
     positives.
 
+    If the optional keyword-only argument ``current_thread_only`` is set to *True*, the
+    plugin will only track memory allocations made by the current thread and all other
+    allocations will be ignored.
+
     .. tip::
 
        You can pass the ``--memray-bin-path`` argument to ``pytest`` to specify

pyproject.toml

@@ -20,7 +20,7 @@ maintainers = [
 requires-python = ">=3.8"
 dependencies = [
   "pytest>=7.2",
-  "memray>=1.5",
+  "memray>=1.12",
 ]
 optional-dependencies.docs = [
   "furo>=2022.12.7",

src/pytest_memray/marks.py

@@ -186,15 +186,27 @@ def _passes_filter(
 def limit_memory(
     limit: str,
     *,
+    current_thread_only: bool = False,
     _result_file: Path,
     _config: Config,
     _test_id: str,
 ) -> _MemoryInfo | _MoreMemoryInfo | None:
     """Limit memory used by the test."""
     reader = FileReader(_result_file)
-    allocations: list[AllocationRecord] = list(
-        reader.get_high_watermark_allocation_records(merge_threads=True)
-    )
+    allocations: list[AllocationRecord] = []
+    if current_thread_only:
+        main_thread = reader.metadata.main_thread_id
+        allocations.extend(
+            record
+            for record in reader.get_high_watermark_allocation_records(
+                merge_threads=False
+            )
+            if record.tid == main_thread
+        )
+    else:
+        allocations.extend(
+            reader.get_high_watermark_allocation_records(merge_threads=True)
+        )
     max_memory = parse_memory_string(limit)
     total_allocated_memory = sum(record.size for record in allocations)
 
@@ -225,14 +237,22 @@ def limit_leaks(
     location_limit: str,
     *,
     filter_fn: Optional[LeaksFilterFunction] = None,
+    current_thread_only: bool = False,
     _result_file: Path,
     _config: Config,
     _test_id: str,
 ) -> _LeakedInfo | None:
     reader = FileReader(_result_file)
-    allocations: list[AllocationRecord] = list(
-        reader.get_leaked_allocation_records(merge_threads=True)
-    )
+    allocations: list[AllocationRecord] = []
+    if current_thread_only:
+        main_thread_id = reader.metadata.main_thread_id
+        allocations.extend(
+            record
+            for record in reader.get_leaked_allocation_records(merge_threads=False)
+            if record.tid == main_thread_id
+        )
+    else:
+        allocations.extend(reader.get_leaked_allocation_records(merge_threads=True))
 
     memory_limit = parse_memory_string(location_limit)
 

tests/test_pytest_memray.py

@@ -868,3 +868,53 @@ def test_memory_alloc_fails():
     )
     result = pytester.runpytest("--memray")
     assert result.ret == ExitCode.OK
+
+
+def test_limit_memory_in_current_thread(pytester: Pytester) -> None:
+    pytester.makepyfile(
+        """
+        import pytest
+        from memray._test import MemoryAllocator
+        allocator = MemoryAllocator()
+        import threading
+        def allocating_func():
+            for _ in range(10):
+                allocator.valloc(1024*5)
+                # No free call here
+
+        @pytest.mark.limit_memory("5KB", current_thread_only=True)
+        def test_memory_alloc_fails():
+            t = threading.Thread(target=allocating_func)
+            t.start()
+            t.join()
+        """
+    )
+
+    result = pytester.runpytest("--memray")
+
+    assert result.ret == ExitCode.OK
+
+
+def test_leaks_in_current_thread(pytester: Pytester) -> None:
+    pytester.makepyfile(
+        """
+        import pytest
+        from memray._test import MemoryAllocator
+        allocator = MemoryAllocator()
+        import threading
+        def allocating_func():
+            for _ in range(10):
+                allocator.valloc(1024*5)
+                # No free call here
+
+        @pytest.mark.limit_leaks("5KB", current_thread_only=True)
+        def test_memory_alloc_fails():
+            t = threading.Thread(target=allocating_func)
+            t.start()
+            t.join()
+        """
+    )
+
+    result = pytester.runpytest("--memray")
+
+    assert result.ret == ExitCode.OK