Import memory profiler

improf.py

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)