Memory & Time Profiling Decorators

profiling.py

"""
################
# Installation #
################
memory_profiler:
    pip install memory-profiler

pympler:
    pip install pympler

line_profiler:
    git clone https://github.com/rkern/line_profiler.git
    find line_profiler -name '*.pyx' -exec cython {} \;
    pip install ./line_profiler


#########
# Usage #
#########
@timeprof
def slow_func(..):
    ...

@memprof
def bloated_func(..):
    ...

@memprof_by_type
def bloated_func(..):
    ...

@memprof_by_line
def bloated_func(..):
    ...

@memprof(by_line=False)
def bloated_func(..):
    ...

@memprof(by_type=False)
def bloated_func(..):
    ...


############
# Overhead #
############
"by_type" memory profiling: High
"by_line" memory profiling: Low
"by_line" time   profiling: Low


#########
# Notes #
#########
Don't stack any of these decorators on top of one another.
For example, don't do any of these:

@timeprof
@memprof
def slow_and_bloated_func(..):
    ...


@timeprof
@memprof_by_line
def slow_and_bloated_func(..):
    ...


@memprof_by_type
@memprof_by_line
def bloated_func(..):
    ...

"""

from functools import wraps
from io import StringIO


def timeprof_by_line(fn):
    from line_profiler import LineProfiler

    @wraps(fn)
    def inner(*a, **kw):
        lp = LineProfiler()
        ret = lp(fn)(*a, **kw)
        lp.print_stats(output_unit=0.001, stripzeros=True)
        return ret

    return inner


timeprof = timeprof_by_line


def memprof_by_type(fn):
    from pympler.tracker import SummaryTracker

    @wraps(fn)
    def inner(*a, **kw):
        tr = SummaryTracker()
        ret = fn(*a, **kw)
        tr.print_diff()
        return ret

    return inner


def memprof_by_line(fn):
    from memory_profiler import profile

    @wraps(fn)
    def inner(*a, **kw):
        bio = StringIO()
        ret = profile(fn, stream=bio)(*a, **kw)
        bio.seek(0)
        print(bio.read())
        return ret

    return inner


def memprof(fn_=None, by_line: bool = True, by_type: bool = True):
    if by_line:
        from memory_profiler import profile

    if by_type:
        from pympler.tracker import SummaryTracker

    def outer(fn):
        @wraps(fn)
        def inner(*a, **kw):
            bio = None
            type_tracker = None
            final_fn = fn

            if by_line:
                bio = StringIO()
                final_fn = profile(fn, stream=bio)

            if by_type:
                type_tracker = SummaryTracker()

            ret = final_fn(*a, **kw)

            if type_tracker is not None:
                type_tracker.print_diff()

            if bio is not None:
                bio.seek(0)
                print(bio.read())

            return ret

        return inner

    if fn_ is not None:
        return outer(fn_)
    else:
        return outer