Kenility - Python - Stopwatch context manager

stopwatch.py

# We create a simple python context manager to log measure time of our code.
from timeit import default_timer as timer

class stopwatch(object):
    def __init__(self, description, log_fn=None):
        self.description = description
        self.log_fn = log_fn
    def __enter__(self):
        self.start = timer()
        return self
    def __exit__(self, t_type, value, traceback):
        took = (timer() - self.start) * 1000
        if self.log_fn:
            self.log_fn("Doing '%s' took %2.2f ms of time" % (self.description, took))

def log_msg(msg):
    print(msg)
    
if __name__ == '__main__':
    import time
    with stopwatch("this task!", log_fn=log_msg):
        time.sleep(1) # code that takes time!
# Output:
# >>> Doing 'this task!' took 1005.40 ms of time