Is there a way to directly "decorate" a block of Python code?

Question

I have a bunch of code:

statement1(args)
statement2(args)
statement3(args)
statement4(args)
statement5(args)

I want to divide the statements into blocks and write to a log after each block. The logging is a bit complex: I want to log things like the running time of each block and the state of a particular data structure after the block executes. So I created a decorator called log_block which handles all of these details. Now my code looks like this:

@log_block()
def block1():
    statement1(args)
    statement2(args)

@log_block()
def block2()
    statement3(args)

@log_block()
def block3():
    statement4(args)
    statement5(args)

block1()
block2()
block3()

This works just fine, but it's a little clunky. It's annoying that I have to separately call the three block functions, and if I want to share a variable between the blocks then I either have to give the block functions arguments and return statements or use global variables, neither of which is particularly palatable. What I really want is syntax which looks like this:

@log_block()
    statement1(args)
    statement2(args)

@log_block()
    statement3(args)

@log_block()
    statement4(args)
    statement5(args)

so that I am decorating the statements directly rather than enclosing them in auxiliary block functions. Is there any way to achieve something like this?

Answer 1

Context managers are exactly what you are looking for. You use them with the with statement, and they define code to be run on entering and exiting the with block.