Implement python decorator to record function's memory usage (or time usage)¶

Sometimes, I want to know the memory usage of several functions and I don't like to write duplicated codes in every related function. Therefore, a general decorator can help.

Install dependency¶

pip install psutil

psutil is a convenient tool to get memory usage of processes. It supports multiple platforms like Windows, macOS and Linux.

Decorator Code¶

import os
from functools import wraps
import psutil


# this decorator is used to record memory usage of the decorated function
def record_mem_usage(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        process = psutil.Process(os.getpid())
        mem_start = process.memory_info()[0]
        rt = func(*args, **kwargs)
        mem_end = process.memory_info()[0]
        diff_KB = (mem_end - mem_start) // 1000
        print('memory usage of %s: %s KB' % (func.__name__, diff_KB))
        return rt
    return wrapper

As we can see in the code, we use psutil.Process(os.getpid()).memory_info()[0] to get the total memory before and after running the decorated function. The difference between the former and the latter is the memory usage that we want to record.

Test the decorator¶

# store is used to hold data, so that python's GC does not affect our test
store = []


@record_mem_usage
def example1():
    lot_of_mem = list(range(10000))
    store.append(lot_of_mem)


@record_mem_usage
def example2():
    lot_of_mem = list(range(20000))
    store.append(lot_of_mem)


# run tests
example1()
example2()

Output the above code looks like:

memory usage of example1: 327 KB
memory usage of example2: 872 KB

What about recording the usage of time?¶

The mechanism to record time usage is similar.

To record time usage, we can utilize time.time():

import time
from functools import wraps


# this decorator is used to record time usage of the decorated function
def record_time_usage(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        time_start = time.time()
        rt = func(*args, **kwargs)
        time_end = time.time()
        diff_ms = (time_end - time_start) * 1000
        print('time usage of %s: %sms' % (func.__name__, diff_ms))
        return rt
    return wrapper

To record CPU time usage, we can utilize time.process_time():

import time
from functools import wraps


# this decorator is used to CPU time usage of the decorated function
def record_cpu_time_usage(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        time_start = time.process_time()
        rt = func(*args, **kwargs)
        time_end = time.process_time()
        diff_ms = (time_end - time_start) * 1000
        print('CPU time usage of %s: %sms' % (func.__name__, diff_ms))
        return rt
    return wrapper