Combining Python's with-syntax and decorators in a single class

 This notebook should show the differences between using a decorator and the with syntax in python. Both will be shown at the example of measuring the time that is required to execute a specific part of code.

With-Syntax

The with syntax is very usefull especially in cases, where ressources are setup and are freed after their use, such as:

1. Opening / Closing files
2. Setting / releasing a lock or semaphore
3. Ensuring a database commit and rolling back on an exception (see example 3 from here)

In order to implement them, there is already a ready to use AbstractContextManager class available, which just requires to implement a __enter__ and __exit__ method to implement e.g. a time keeper decorator:

In [1]:

from contextlib import AbstractContextManager

import time

class TrackTime(AbstractContextManager):
    def __init__(self):
        pass
    
    def __enter__(self):
        self.start_time = time.time()
        
    def __exit__(self, *args):
        end_time = time.time()
        print(f"Total running time was: {end_time - self.start_time}")

In [2]:

with TrackTime():
    time.sleep(2)

Total running time was: 2.0008251667022705

A decorator can be used directly to annotate a function and achieve a similar result like above: Tracking the time of that function with a decorator:

In [3]:

def tracktimefunction(func, *args, **kwargs):
    def wrapping_function(*args, **kwargs):
        start_time = time.time()

        func(*args, **kwargs)

        end_time = time.time()
        print(f"Total running time was: {end_time - start_time}")
        return
    return wrapping_function

In [4]:

@tracktimefunction
def wait_2s():
    time.sleep(2)

wait_2s()

Total running time was: 2.0113117694854736

The decorator can be easily integrated into the same class from above that served the purpose to time the time of a specific section of the code with the with syntax. This can be achieved by making the class "callable". The method __call__ will be responsible now for wrapping the function internally and it can be further simplified to use the with syntax with the class itself:

In [5]:

from contextlib import AbstractContextManager

import time

class TrackTime(AbstractContextManager):
    def __init__(self):
        pass
    
    def __enter__(self):
        self.start_time = time.time()
        
    def __exit__(self, *args):
        end_time = time.time()
        print(f"Total running time was: {end_time - self.start_time}")

    def __call__(self, func, *args, **kwargs):
        def wrapping_function(*args, **kwargs):
            with self:
                return func(*args, **kwargs)
        return wrapping_function

Now both, the decorator as well as the with syntax is possible within the same class:

In [6]:

@TrackTime()
def wait_2s():
    time.sleep(2)

wait_2s()

with TrackTime():
    time.sleep(3)

Total running time was: 2.0006048679351807
Total running time was: 3.0028138160705566