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

BioBroo 0
Viewed 40

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)
with-statement pytorch python
2 Answers

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"
chepner0

I was wondering if it could be accomplished with:

def myrunner(myfunc,usecontext=True):
    if usecontext:
        with blahblah():
             myfunc()
    else:
        myfunc()

def myfunc()...

myrunner(myfunc,usecontext=false)

But looking at it, I don't see how the context can effect myfunc.

RufusVS0
Related