How to use the command "with" conditionally without duplication of code

Question

I am trying to use with before executing a block of code, but only if a condition is met, but (at least the common) usage of with doesn't appear to support that unless I duplicate the block of code.

More concretely, I know I can do the following:

if condition:
    with blah_blah():
        my_code_block
else:
    my_code_block

But that's unsatisfying during development since any change I make to my_code_block must be made twice. What I want to do (conceptually) is:

if condition:
    with blah_blah():
else:
    my_code_block

That doesn't work, though. Is there a way to accomplish what I'm trying to do?

For anyone that's interested in my particular use, I'm trying to write code that runs a batch of examples in pytorch, with torch.no_grad() if I'm in evaluation mode and without it if I'm in train mode. So what I want to do becomes

if mode == 'eval':
    with torch.no_grad():
else:
    run_batch(features, labels)

Answer 1

Use the with statement, but with a nullcontext context manager if necessary.

from contextlib import nullcontext

with blah_blah() if condition else nullcontext():
    my_code_block

nullcontext takes an optional argument that it will return if your with statement expects something to be bound with as. For example,

with nullcontext("hello") as f:
    print(f)  # outputs "hello"