Trace peak memory usage in python

trace_peak.py

from contextlib import contextmanager
import tracemalloc

@contextmanager
def mtrace():
    tracemalloc.start()
    yield
    current, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    print(f"Peak memory usage: {peak / 1024 ** 2:.2f} MiB")
    
    
# usage
with mtrace():
    func()