benchmark_class.py

import time

class Benchmark():

    def __init__(self, txt):
        self.txt = str(txt)
        self.start = time.time()
 
    def __enter__(self):
        return self
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        end = time.time()
        print(self.txt + " took " + str(end - self.start) + " seconds.")


def example_function():
    time.sleep(5)
 
 
if __name__ == '__main__':
    with Benchmark("example_function"):
        example_function()

benchmark_function.py

import time

def benchmark(func):
    def _timer(*args, **kwargs):
        start = time.time()
        value = func(*args, **kwargs)
        end = time.time()
        print(str(func.__name__) + " took " + str(end - start) + " seconds.")
        return value
    return _timer