Source

Dan Bader, Python Tricks: The Book, 2.3 Context Managers and the with statement

Description

Implement a context manager that measures the execution time of a code block using the time.time() function. Write two versions:

  1. a class-based version
  2. a generator-based version used with a contextmanager decorator

My solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

from contextlib import contextmanager
from time import time


class BlockTimer:
    """Class-based version.
    
    Define `__enter__()` and `__exit__()`, which are invoked
    when entering and exiting the context established by the `with`
    keyword.
    
    Also define `__repr__()` to return the string representation used
    by `print()`.
    """
    def __init__(self):
        self.begin = 0
        self.end = 0
        
    def __enter__(self):
        self.begin = time()
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.begin:
            self.end = time()
            print(self)

    def __repr__(self):
        return f'Code block executed in {self.end - self.begin :.02f}s'


@contextmanager
def block_timer():
    """Generator-based version.
    
    `yield` to return to context.
    """
    try:
        begin = time()
        yield
    finally:
        print(f'Code block executed in {time() - begin :.02f}s') 
    


if __name__ == '__main__':
    from time import sleep

    with BlockTimer() as bt:
        sleep(3)

    with block_timer() as bt:
        sleep(3)