Detect Memory Leaks in Python with Weak References

import gc
import sys
import weakref
import time
from collections import defaultdict

class LeakDetector:
    def __init__(self):
        self._tracked = defaultdict(list)

    def track_class(self, cls):
        """Track all instances of a class for leak detection."""
        old_init = cls.__init__
        def new_init(self, *args, **kwargs):
            old_init(self, *args, **kwargs)
            weak_self = weakref.ref(self)
            # Store creation traceback
            import traceback
            try:
                raise RuntimeError
            except RuntimeError:
                trace = traceback.extract_stack()[:-2]
            LeakDetector._current_instance()._tracked[cls.__name__].append((weak_self, trace))
        cls.__init__ = new_init

    @staticmethod
    def _current_instance():
        if not hasattr(LeakDetector, '_instance'):
            LeakDetector._instance = LeakDetector()
        return LeakDetector._instance

    @staticmethod
    def check_for_leaks():
        detector = LeakDetector._current_instance()
        leaks_found = 0
        for class_name, instances in detector._tracked.items():
            for weak_ref, trace in instances:
                obj = weak_ref()
                if obj is not None and sys.getrefcount(obj) > 2:
                    if leaks_found == 0:
                        print(f"\n{'='*60}")
                        print("LEAK DETECTED!")
                        print("="*60)
                    print(f"\nLeaked {class_name} instance at {id(obj):#x}")
                    print("Creation traceback:")
                    for filename, lineno, func, text in trace:
                        print(f"  File '{filename}', line {lineno}, in {func}")
                        if text:
                            print(f"    {text}")
                    leaks_found += 1
        if leaks_found == 0:
            print("No memory leaks detected.")
        return leaks_found

def demonstrate_leak_detection():
    # Example: A class that will leak if not cleaned up
    class LeakyResource:
        def __init__(self, name):
            self.name = name
            # Simulate a circular reference that prevents garbage collection
            self._leak = self  # Intentionally create circular reference

    class CleanResource:
        def __init__(self, name):
            self.name = name
            self._data = f"Data for {name}"

    # Enable leak tracking
    LeakDetector.track_class(LeakyResource)
    LeakDetector.track_class(CleanResource)

    # Create instances - some will be properly cleaned, some will leak
    def create_leaked_resources():
        for i in range(3):
            resource = CleanResource(f"clean_{i}")  # Will be collected
            _ = LeakyResource(f"leaky_{i}")  # Circular ref prevents collection

    create_leaked_resources()
    
    # Force garbage collection
    gc.collect()
    
    # Detect leaks
    LeakDetector.check_for_leaks()

if __name__ == "__main__":
    print("Memory Leak Detection Demo")
    print("Creating resources and checking for leaks...")
    demonstrate_leak_detection()