guppyでのヒープ測定をwithとかデコレータで出来るようにするスニペット

measure_heap.py

# -*- coding: utf-8 -*-

from __future__ import print_function

from guppy import hpy


class measure_heap(object):
    '''
    メモリ消費量を測定する。
    Warning:
        guppyというモジュールが必要

    Usage:
        with measure_heap() as m:
            print('in with')

        @measure_heap
        def hoge():
            print('in dec')
        hoge()
    '''

    def __init__(self, func=None, printer=None):
        self.func = func
        self.printer = printer

    def _pre_measurement(self):
        return hpy()


    def __call__(self):
        if not self.func:
            raise ValueError('this object should be used as decorator or context holder')

        h = self._pre_measurement()
        self.func()
        measurement = h.heap
        if self.printer:
            self.printer(measurement())
        else:
            print(measurement())

    def __enter__(self, printer=None):
        self.printer = printer
        self.h = self._pre_measurement()
        return self

    def __exit__(self, type, value, traceback):
        measurement = self.h.heap
        if self.printer:
            self.printer(measurement())
        else:
            print(measurement())