Python context manager that measures time

Question

I am struggling to make a piece of code that allows to measure time spent within a "with" statement and assigns the time measured (a float) to the variable provided in the "with" statement.

import time

class catchtime:
    def __enter__(self):
        self.t = time.clock()
        return 1

    def __exit__(self, type, value, traceback):
        return time.clock() - self.t

with catchtime() as t:
    pass

This code leaves t=1 and not the difference between clock() calls. How to approach this problem? I need a way to assign a new value from within the exit method.

PEP 343 describes in more detail how contect manager works but I do not understand most of it.

Answer 1

The top rated answer can give the incorrect time

As noted already by @Mercury, the top answer by @Vlad Bezden, while slick, is technically incorrect since the value yielded by t() is also potentially affected by code executed outside of the with statement. For example, if you execute time.sleep(5) after the with statement but before the print statement, then calling t() in the print statement will give you ~6 sec, not ~1 sec.

This can be avoided by doing the print command inside the context manager as below:

from time import perf_counter
from contextlib import contextmanager


@contextmanager
def catchtime() -> float:
    start = perf_counter()
    yield lambda: perf_counter() - start
    # Note: print is included here to guarantee only time of code inside the CM is measured
    print(f'Time: {perf_counter() - start:.3f} seconds')

Notice how sleep(5) causes the incorrect time to be printed:

from time import sleep

with catchtime() as t:
    sleep(1)

# >>> "Time: 1.000 seconds"

sleep(5)
print(f'Time: {t():.3f} seconds')

# >>> "Time: 6.000 seconds"

A more robust solution (recommended)

This code is similar to the excellent answer given by @BrenBarn execept that it:

1. Automatically prints the executed time as a formatted string (remove this by deleting the final print(self.readout))
2. Saves the formatted string for later use (self.readout)
3. Saves the float result for later use (self.time)

from time import perf_counter


class catchtime:
    def __enter__(self):
        self.time = perf_counter()
        return self

    def __exit__(self, type, value, traceback):
        self.time = perf_counter() - self.time
        self.readout = f'Time: {self.time:.3f} seconds'
        print(self.readout)

Notice how the intermediate sleep(5) commands no longer have an effect on the printed time.

from time import sleep

with catchtime() as t:
    sleep(1)

# >>> "Time: 1.000 seconds"

sleep(5)
print(t.time)

# >>> 1.000283900000009

sleep(5)
print(t.readout)

# >>> "Time: 1.000 seconds"