import tracemalloc
import time

class MyTraceMalloc:
    def __init__(self):
        tracemalloc.start(25)  # 保留25个堆栈级别
        self.before_snapshot = tracemalloc.take_snapshot()
        self.current_snapshot = None
        self.top_stats = []

    # 内存跟踪情况输出
    def display_top_stats(self):
        self.current_snapshot = tracemalloc.take_snapshot()
        self.top_stats = self.current_snapshot .compare_to(self.before_snapshot, 'lineno')

        print("[ Top 10 differences ]")
        for stat in self.top_stats[:10]:
            print(stat)
        current, peak = tracemalloc.get_traced_memory()
        print(f"当前内存分配: {current / 1024 / 1024:.2f} MiB")
        print(f"峰值内存分配: {peak / 1024 / 1024:.2f} MiB")

        self.before_snapshot = self.current_snapshot

    def trace_memory(self):
        while True:
            time.sleep(60)  # 每隔60秒输出一次
            self.display_top_stats()