Execution timing with context manager

Question

I was looking for a solution for timing code execution. Of course, I found all sorts of solutions mostly suggesting the use of timeit module - which is pretty cool - or just using the time module. Neither of those really provides what I'm after.

I wanted a simple solution that

• I can quickly add to the code without changing the structure or logic
• can cover a single line or code block
• elegant and clean enough to even leave it in

I think I finally figured it out - using the context manager's dunder methods __enter__ and __exit__.

The idea is we capture the time when the context manager starts (__enter__), then just let the code run in the context-manager-block, finally, we capture the time when the context manager ends (__exit__) and do something with the result (print and/or return).

So here's the snippet:

import time


class ExecutionTimer:

    def __init__(self, message='Execution time', unit='s'):
        self.message = message
        self.unit = unit
        self.start_time = None
        self.end_time = None
        self.result = None

    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        self.result = self.end_time - self.start_time
        self.print_time()

    def print_time(self):
        elapsed_time = self.result
        if self.unit == 'ms':
            elapsed_time *= 1000  # convert to milliseconds
        elif self.unit == 'us':
            elapsed_time *= 1e6  # convert to microseconds
        elif self.unit == 'ns':
            elapsed_time *= 1e9  # convert to nanoseconds

        print(f"{self.message}: {elapsed_time}{self.unit}")


if __name__ == '__main__':
    
    start_time = time.time()
    time.sleep(1)
    end_time = time.time()
    print(f"Execution (s): {end_time - start_time}s")

    with ExecutionTimer(message='Execution (s)', unit='s'):
        time.sleep(1)

    with ExecutionTimer(message="Execution (ms)", unit='ms'):
        time.sleep(1)

    with ExecutionTimer(message="Execution (us)", unit='us'):
        time.sleep(1)

    with ExecutionTimer(message="Execution (ns)", unit='ns'):
        time.sleep(1)

    # or just capture the value
    with ExecutionTimer(message="Execution (ns)", unit='ns') as my_timer:
        time.sleep(1)

    # notice we are not in the context manager any more
    print(my_timer.result)
    print(my_timer.result)

and its output:

Execution (s): 1.0000789165496826s
Execution (s): 1.0000693798065186s
Execution (ms): 1000.067949295044ms
Execution (us): 1000072.0024108887us
Execution (ns): 1000071287.1551514ns
Execution (ns): 1000077962.8753662ns
1.0000779628753662
1.0000779628753662

Process finished with exit code 0

There's definitely some room for improvement. It can be integrated with logging to only emit messages to a certain log level OR it could use a "dry run" parameter to control the execution print method. Etc.

Feel free to use, test, and modify it as you please. It would be awesome if you could leave your constructive comments, ideas, and improvements so we can all learn from them.

Cheers!