memory_tracer.py

import logging
import tracemalloc

tracemalloc.start(10)

logger = logging.getLogger(__name__)


class MemoryTracer:
    def __init__(self):
        self.snapshots = []

    def capture_snapshot(self):
        snapshot = tracemalloc.take_snapshot()
        self.snapshots.append(snapshot)

    def log_top_n_size_increases_since_last_snapshot(self, n):
        def by_size_diff(statistic_diff):
            return statistic_diff.size_diff

        if len(self.snapshots) < 2:
            logger.info(
                "Requires at least 2 snapshots "
                "in order to log size increases, continuing."
            )
            return

        logger.info("Logging biggest object size increases since last snapshot")

        stats = self.snapshots[-1].compare_to(self.snapshots[-2], "filename")

        top_n_stats_by_size_diff = sorted(stats, key=by_size_diff, reverse=True)[:n]
        logger.info("\n".join([str(obj) for obj in top_n_stats_by_size_diff]))

    def log_largest_n_objects(self, n):
        def by_size(statistic):
            return statistic.size

        if not len(self.snapshots):
            logger.info(
                "At least one snapshot required to log largest snapshots, continuing."
            )
            return

        logger.info("Logging largest objects currently in memory")

        snapshot = self.snapshots[-1]
        statistics = snapshot.statistics("filename")

        largest_objects = sorted(statistics, key=by_size, reverse=True)[:n]
        logger.info("\n".join([str(obj) for obj in largest_objects]))


# Create a memory tracer (maybe as a singleton)
memory_tracer = MemoryTracer()

# After some computation
memory_tracer.capture_snapshot()

# After some more computation
memory_tracer.capture_snapshot()

memory_tracer.log_top_n_size_increases_since_last_snapshot(10)
memory_tracer.log_largest_n_objects(10)