A contextmanager for pausing garbage collection.

pause_garbage_collection.py

"""
pause_garbage_collection

Creates a contextmanager to pause garbage collection.
Garbage may still be collected explicitly via `gc.collect()`

with pause_garbage_collection():
    foo()
    bar()
"""
from contextlib import contextmanager
from typing import Iterator
import gc


@contextmanager
def pause_garbage_collection(force_collect: bool = False) -> Iterator[None]:
    """Pause garbage collection while a task is run in this context manager."""
    try:
        gc.disable()
        yield
    finally:
        if force_collect:
            gc.collect()
        gc.enable()


class _MemoryLeakObject():
    """Create a memory leak via circular references"""
    def __init__(self):
        self.ref = None

    def __del__(self):
        pass

    def leak(self):
        """Create a circular reference, making cleaning this up GC's job."""
        self.ref = self

    def plug(self):
        """Remove the circular reference, allowing ref counting to do the job."""
        self.ref = None


def _test_pgc_contextmanager():
    """Test the pause garbage collection context manager."""
    with pause_garbage_collection(True):
        for i in range(30000000):
            mlo = _MemoryLeakObject()
            mlo.leak()
            if (i+1) % 10000000 == 0:
                print("Currently tracking {} objects.".format(len(gc.get_objects())))
                print("Collecting garbage...")
                gc.collect()
                print("{} objects remain.".format(len(gc.get_objects())))
                print("Collection complete!\n")


def _test_pgc_noleak():
    """Test that _MemoryLeakObject's with plugged leaks are reclaimed by the reference counter."""
    with pause_garbage_collection(True):
        for i in range(30000000):
            mlo = _MemoryLeakObject()
            mlo.leak()
            mlo.plug()
            if (i+1) % 10000000 == 0:
                print("Currently tracking {} objects.".format(len(gc.get_objects())))
                print("Collecting garbage...")


if __name__ == "__main__":
    _test_pgc_contextmanager()
    _test_pgc_noleak()