profiler.py

import time
from contextlib import contextmanager
from functools import wraps
from collections import defaultdict

class Profiler:
    def __init__(self):
        self.timings = defaultdict(list)

    @contextmanager
    def block(self, block_name):
        start = time.time()
        try:
            yield
        finally: # <- whatever happens during block execution I can record the timing
            end = time.time()
            self.timings[block_name].append(end - start)

    def wrap(self, func):
        @wraps(func) # <- always remember to wrap a function when decorating it
        def wrapper(*args, **kwargs):
            with self.block(func.__name__):
                return func(*args, **kwargs)
            return wrapper

    def report(self):
        """
        Print number of calls, min/max/avg/total time of each profiled function or block
        """
        for name, timing in self.timings.items():
            t_max = max(timing)
            t_min = min(timing)
            t_sum = sum(timing)
            t_avg = t_sum/len(timing)
            print(f"{name}: max: {t_max}s, min: {t_min}s, total: {t_sum}s, avg: {t_avg}s")
            
# module instance that can be shared as one single profiler with decorator          
profiler = Profiler()