"""
This module contains functions dealing with assessing code performance.
"""
import time
from datetime import timedelta
import psutil
__all__ = [
'Timer',
'print_efficiency',
]
[docs]class Timer(object):
"""
Implements a context-manager timer.
Args:
ndecimals (int): number of decimals to print (for the seconds).
By default rounds to 0 decimals, i.e., whole seconds.
Examples:
>>> with Timer(): # doctest: +SKIP
... s = [x**2 for x in range(10000000)]
Elapsed time: 0:00:02.172509
"""
def __init__(self, prefix="", ndecimals=None):
self.ndecimals = ndecimals
self.prefix = prefix
def __enter__(self):
self.start = time.time()
def __exit__(self, type, value, traceback):
self.end = time.time()
dt = self.end - self.start
if self.ndecimals is not None:
dt = round(dt, self.ndecimals)
print("{}\n\tElapsed time: {}".format(self.prefix, timedelta(seconds=dt)))
[docs]def print_efficiency(f):
"""
Decorator that print elasped time and change in available memory when calling the decorated function.
"""
def wrapper(*args, **kwargs):
# Start time and memory available.
t0 = time.time()
mem0 = psutil.virtual_memory().available
# Call function.
out = f(*args, **kwargs)
# End time and memomry available.
t1 = time.time()
mem1 = psutil.virtual_memory().available
# Print efficiency.
print('Efficiency function "{}":\n'.format(f.__name__) +
'\tTime elapsed: {:.4f} s\n'.format(t1 - t0) +
'\tChange available memory: {:.2f} MB'.format((mem1 - mem0)/1e6))
return out
return wrapper