Skip to content

Instantly share code, notes, and snippets.

@baverman

baverman/improf.py

Created Jun 10, 2019
Embed
What would you like to do?
Import memory profiler
import sys
from importlib._bootstrap import _find_and_load
from resource import getrusage, RUSAGE_SELF
from contextlib import contextmanager
IMPORT_ENTRY = _find_and_load.__code__
def get_mem():
return getrusage(RUSAGE_SELF).ru_maxrss
@contextmanager
def profile(stream=None):
p = ImportMemoryProfiler()
p.start()
try:
yield
finally:
p.stop()
p.summary(stream)
p.report(stream)
class ImportMemoryProfiler:
def __init__(self):
self.top = ['__top__', get_mem(), []]
self.stack = [self.top]
self.current = self.top
def trace(self, frame, event, ctx):
if event == 'call':
if frame.f_code is IMPORT_ENTRY:
e = [frame.f_locals['name'], get_mem(), []]
self.stack.append(e)
self.current[2].append(e)
self.current = e
elif event == 'return':
if frame.f_code is IMPORT_ENTRY:
e = self.stack.pop()
e[1] = get_mem() - e[1]
self.current = self.stack[-1]
return self.trace
def start(self):
sys.settrace(self.trace)
def stop(self):
sys.settrace(None)
self.top[1] = get_mem() - self.top[1]
def report(self, stream=None):
def _print(node, level):
child_total = sum(it[1] for it in node[2])
self_usage = node[1] - child_total
print(' '*level + '{} {} {}'.format(node[0], node[1], self_usage), file=stream)
for it in node[2]:
_print(it, level+1)
_print(self.top, 0)
def summary(self, stream=None):
from collections import Counter
stats = Counter()
def _stats(node):
child_total = sum(it[1] for it in node[2])
self_usage = node[1] - child_total
stats[node[0].partition('.')[0]] += self_usage
for it in node[2]:
_stats(it)
_stats(self.top)
print(sum(stats.values()), 'TOTAL', sep='\t', file=stream)
for it, size in stats.most_common(100):
print(size, it, sep='\t', file=stream)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.