Simple context manager for timing Python code

timethis.py

#!/usr/bin/env python
from contextlib import contextmanager
import time


@contextmanager
def timethis(timer=time.perf_counter):
    """
    Small timing function. It is meant to be used as a
    context manager, and it returns a lambda that when
    called returns the time it took for the block to
    execute.

    >>> from time import sleep
    >>> with timethis() as timer:
    ...    sleep(4.5)
    ...
    >>> print('Finished in %.4f seconds.' % timer())
    Finished in 4.5017 seconds.
    """
    start = timer()
    try:
        yield lambda: end - start
    finally:
        end = timer()


if __name__ == '__main__':
    with timethis() as timer:
        for i in range(20):
            print(i)

    print('Operation took', timer(), 'seconds.')