Add a new marker to check for memory leaks

Users have indicated that it will be very useful if the plugin exposes a
way to detect memory leaks in tests. This is possible, but is a bit
tricky as the interpreter can allocate memory for internal caches, as
well as user functions.

To make this more reliable, the new marker will take two parameters:

* The watermark of memory to ignore. If the memory leaked by the test is
  higher than this value, the test will fail and it will pass otherwise.

* The number of warmup runs. This allows to run the test multiple times
  (assuming it passes) before actually checking for leaks. This allows
  to warmup user and interpreter caches.

Signed-off-by: Pablo Galindo <[email protected]>

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,83 @@ 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(max_leaked_memory: str , *, warmups: int=0)
+
+      Fail the execution of the test if the test leaks more memory than allowed.
+
+.. important:: 
+
+      To detect leaks, Memray needs to intercept calls to the Python allocators.
+      This is adds significant overhead, and will slow your test down.
+      Additionally, if the ``warmups`` argument is provided, the test will be run
+      multiple times before the final execution, which will also increase the
+      total time it takes to complete the test run.
+
+When this marker is applied to a test, it will cause the test to fail if the execution
+of the test leaks more memory than allowed. It takes a single positional argument with a
+string indicating the maximum memory 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.
+
+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 ``warmups``. This argument
+indicates the number of times **the test body (and not any of the fixtures)**
+will be executed before the memory leak check is performed. This is useful to
+avoid false positives where memory allocated below the test is cached and
+survives the test run, as well as to account for memory allocated (and not
+released) by the interpreter under normal execution.
+
+.. 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", warmups=3)
+    def test_foobar():
+        pass # do some stuff that cannot leak memory
+
+.. warning::
+
+   Is **very** challenging to write tests that do not "leak" memory in some way.
+   This marker is intended to be used to detect leaks in very simple unit tests
+   and can be challenging to use reliably when applied to integration tests or
+   more complex scenarios. Also, ``pytest`` itself does things that will be interpreted
+   as leaking memory, for example:
+
+   * ``pytest`` captures stdout/stderr by default and will not release the
+     memory associated to these strings until the test ends. This will be
+     interpreted as a leak.
+   * ``pytest`` captures logging records by default and will not release the
+     memory associated to these strings until the test ends. This will be
+     interpreted as a leak.
+   * Memory allocated **in the test body** from fixtures that is only
+     released at fixture teardown time will be interpreted as a leak
+     because the plugin cannot see the fixtures being deallocated.
+   * The Python interpreter has some internal caches that will not be cleared
+     when the test ends. The plugin does is best to warmup all available
+     interpreter caches but there are some that cannot be correctly detected so
+     you may need to allow some small amount of leaked memory or use the
+     ``warmups`` argument to avoid 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
+     can all change from one Python version to another.
+
+   For this reason, using this marker with "0B" as the maximum allowed leaked memory
+   in tests that are not very simple unit tests will almost always require a
+   nonzero value for the ``warmups`` argument and some times that will not be enough.

src/pytest_memray/cpython_caches.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import logging
+import contextlib
+import sys
+import warnings
+from inspect import isabstract
+from typing import Any
+from typing import Generator
+from typing import Set
+from typing import Tuple
+
+try:
+    from _abc import _get_dump
+except ImportError:
+    import weakref
+
+    def _get_dump(cls: Any) -> Tuple[Set[Any], Any, Any, Any]:
+        # Reimplement _get_dump() for pure-Python implementation of
+        # the abc module (Lib/_py_abc.py)
+        registry_weakrefs = set(weakref.ref(obj) for obj in cls._abc_registry)
+        return (
+            registry_weakrefs,
+            cls._abc_cache,
+            cls._abc_negative_cache,
+            cls._abc_negative_cache_version,
+        )
+
+
+@contextlib.contextmanager
+def warmup_bytecode_cache() -> Generator[None, None, None]:
+    prev_prov = sys.getprofile()
+    sys.setprofile(lambda *args: args)
+    try:
+        yield
+    finally:
+        sys.setprofile(prev_prov)
+
+
+def clean_interpreter_caches() -> None:
+    # Clear the warnings registry, so they can be displayed again
+    for mod in sys.modules.values():
+        if hasattr(mod, "__warningregistry__"):
+            del mod.__warningregistry__
+
+    # Flush standard output, so that buffered data is sent to the OS and
+    # associated Python objects are reclaimed.
+    for stream in (sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__):
+        if stream is not None:
+            stream.flush()
+
+    with contextlib.suppress(KeyError):
+        re = sys.modules["re"]
+        re.purge()
+
+    with contextlib.suppress(KeyError):
+        _strptime = sys.modules["_strptime"]
+        _strptime._regex_cache.clear()
+
+    with contextlib.suppress(KeyError):
+        urllib_parse = sys.modules["urllib.parse"]
+        urllib_parse.clear_cache()
+
+    with contextlib.suppress(KeyError):
+        urllib_request = sys.modules["urllib.request"]
+        urllib_request.urlcleanup()
+
+    with contextlib.suppress(KeyError):
+        linecache = sys.modules["linecache"]
+        linecache.clearcache()
+
+    with contextlib.suppress(KeyError):
+        mimetypes = sys.modules["mimetypes"]
+        mimetypes._default_mime_types()
+
+    with contextlib.suppress(KeyError):
+        filecmp = sys.modules["filecmp"]
+        filecmp._cache.clear()
+
+    with contextlib.suppress(KeyError):
+        struct = sys.modules["struct"]
+        struct._clearcache()
+
+    with contextlib.suppress(KeyError):
+        doctest = sys.modules["doctest"]
+        doctest.master = None  # type: ignore
+
+    with contextlib.suppress(KeyError):
+        ctypes = sys.modules["ctypes"]
+        ctypes._reset_cache()
+
+    with contextlib.suppress(KeyError):
+        typing = sys.modules["typing"]
+        for cleanup_fn in typing._cleanups:
+            cleanup_fn()
+
+    with contextlib.suppress(KeyError, AttributeError):
+        fractions = sys.modules["fractions"]
+        fractions._hash_algorithm.cache_clear()
+
+    sys._clear_type_cache()
+
+
+def warm_caches() -> None:
+
+    # Run the clean_interpreter_caches() function under
+    # a dummy tracer to warmup the bytecode cache so it
+    # doesn't interfere when called under tracing.
+    with warmup_bytecode_cache():
+        clean_interpreter_caches()
+
+    # char cache
+    s = bytes(range(256))
+    for i in range(256):
+        s[i : i + 1]  # noqa
+    # unicode cache
+    [chr(i) for i in range(256)]
+    # int cache
+    list(range(-5, 257))
+
+
+@contextlib.contextmanager
+def interpreter_cache_context() -> Generator[None, None, None]:
+    import collections.abc
+    import copyreg
+
+    fs = warnings.filters[:]
+    ps = copyreg.dispatch_table.copy()
+    pic = sys.path_importer_cache.copy()
+    try:
+        import zipimport
+    except ImportError:
+        zdc = None  # Run unmodified on platforms without zipimport support
+    else:
+        zdc = zipimport._zip_directory_cache.copy()  # type: ignore
+    abcs = {}
+    for abc in [getattr(collections.abc, a) for a in collections.abc.__all__]:
+        if not isabstract(abc):
+            continue
+        for obj in abc.__subclasses__() + [abc]:
+            abcs[obj] = _get_dump(obj)[0]
+
+    warm_caches()
+
+    yield
+
+    # Restore some original values.
+    warnings.filters[:] = fs  # type: ignore
+    copyreg.dispatch_table.clear()
+    copyreg.dispatch_table.update(ps)
+    sys.path_importer_cache.clear()
+    sys.path_importer_cache.update(pic)
+    try:
+        import zipimport
+    except ImportError:
+        pass  # Run unmodified on platforms without zipimport support
+    else:
+        zipimport._zip_directory_cache.clear()  # type: ignore
+        zipimport._zip_directory_cache.update(zdc)  # type: ignore
+
+    # Clear ABC registries, restoring previously saved ABC registries.
+    abs_classes = [getattr(collections.abc, a) for a in collections.abc.__all__]
+    for abc in filter(isabstract, abs_classes):
+        for obj in abc.__subclasses__() + [abc]:
+            for ref in abcs.get(obj, set()):
+                if ref() is not None:
+                    obj.register(ref())
+            obj._abc_caches_clear()
+
+
+@contextlib.contextmanager
+def all_logging_disabled(highest_level=logging.CRITICAL):
+    """
+    A context manager that will prevent any logging messages
+    triggered during the body from being processed.
+    :param highest_level: the maximum logging level in use.
+      This would only need to be changed if a custom level greater than CRITICAL
+      is defined.
+    """
+    # two kind-of hacks here:
+    #    * can't get the highest logging level in effect => delegate to the user
+    #    * can't get the current module-level override => use an undocumented
+    #       (but non-private!) interface
+
+    previous_level = logging.root.manager.disable
+
+    logging.disable(highest_level)
+
+    try:
+        yield
+    finally:
+        logging.disable(previous_level)