Skip to content
Advertisement

Function wrapper for multiple code lines in Python

I would like to wrap some code around a function (e.g. print timing information) and reuse this function for multiple code blocks:

start = time.time()
print('Started adding numbers at {}'.format(start))
a = 1
b = 2
c = a + b
end = time.time()
print('Finished adding numbers in {} seconds'.format(end - start)')

Is it possible to wrap multiple lines of code into a wrapper function, without defining functions for each code blocks? One way would be to define a function and use a decorator but I would like to avoid doing this for every code block:

@print_time
def foo():
    a = 1
    b = 2

    return a + b

Advertisement

Answer

You could create a context manager class and execute your code in a with block. This would make for a compact and highly reusable timer function. Something like this:

import time

class TimeClass():
    def __init__(self):
        pass
    def __enter__(self):
        self.start = time.time()
    def __exit__(self, type, value, traceback):
        self.end = time.time()
        print('Finished executing in {} seconds'.format(self.end - self.start))

You could use the class like this:

with TimeClass():
    # your code
    pass

Check this link out for more details.

User contributions licensed under: CC BY-SA
2 People found this is helpful
Advertisement