Python context manager: conditionally executing body?

Question

I'm writing an MPI-based application (but MPI doesn't matter in my question, I mention it only to expose the rationale) and in some cases, when there is less work items than processes, I need to create a new communicator excluding the processes that have nothing to do. Finally, the new communicator has to be freed by the processes that have work to do (and only by them).

A neat way to do that would be to write:

with filter_comm(comm, nworkitems) as newcomm:
    ... do work with communicator newcomm...

the body being executed only by the processes that have work to do.

Is there a way in a context manager to avoid executing the body? I understand that context managers have rightfully been designed to avoid hiding control flows, but I wonder if it is possible to circumvent that, since in my case I think it would be justified for clarity sake.

Answer 1

The ability to conditionally skip context manager body has been proposed but rejected as documented in PEP 377.

I did some research about alternatives. Here are my findings.

First let me explain the background of my code examples. You have a bunch of devices you want to work with. For every device you have to acquire the driver for the device; then work with the device using the driver; and lastly release the driver so others can acquire the driver and work with the device.

Nothing out of the ordinary here. The code looks roughly like this:

driver = getdriver(devicename)
try:
  dowork(driver)
finally:
  releasedriver(driver)

But once every full moon when the planets are not aligned correctly the acquired driver for a device is bad and no work can be done with the device. This is no big deal. Just skip the device this round and try again next round. Usually the driver is good then. But even a bad driver needs to be released otherwise no new driver can be acquired.

(the firmware is proprietary and the vendor is reluctant to fix or even acknowledge this bug)

The code now looks like this:

driver = getdriver(devicename)
try:
  if isgooddriver(driver):
    dowork(driver)
  else:
    pass # do nothing or log the error
finally:
  release(driver)

That is a lot of boilerplate code that needs to be repeated everytime work needs to be done with a device. A prime candidate for python's context manager also known as with statement. It might look like this:

# note: this code example does not work
@contextlib.contextmanager
def contextgetdriver(devicename):
  driver = getdriver(devicename)
  try:
    if isgooddriver(driver):
      yield driver
    else:
      pass # do nothing or log the error
  finally:
    release(driver)

And then the code when working with a device is short and sweet:

# note: this code example does not work
with contextgetdriver(devicename) as driver:
  dowork(driver)

But this does not work. Because a context manager has to yield. It may not not yield. Not yielding will result in a RuntimeException raised by contextmanager.

So we have to pull out the check from the context manager

@contextlib.contextmanager
def contextgetdriver(devicename):
  driver = getdriver(devicename)
  try:
    yield driver
  finally:
    release(driver)

and put it in the body of the with statement

with contextgetdriver(devicename) as driver:
  if isgooddriver(driver):
    dowork(driver)
  else:
    pass # do nothing or log the error

This is ugly because now we again have some boilerplate that needs to be repeated everytime we want to work with a device.

So we want a context manager that can conditionaly execute the body. But we have none because PEP 377 (suggesting exactly this feature) was rejected. Thanks for nothing guido.

(actually thank you guido for the beautiful and powerful python language but i find this particular decision questionable)

I found that abusing a generator works quite well as a replacement of a context manager which can skip the body

def generatorgetdriver(devicename):
  driver = getdriver(devicename)
  try:
    if isgooddriver(driver):
      yield driver
    else:
      pass # do nothing or log the error
  finally:
    release(driver)

But then the calling code looks very much like a loop

for driver in generatorgetdriver(devicename):
  dowork(driver)

If you can live with this (please don't) then you have a context manager that can conditionaly execute the body.

It seems hat the only way to prevent the boilerplate code is with a callback

def workwithdevice(devicename, callback):
  driver = getdriver(devicename)
  try:
    if isgooddriver(driver):
      callback(driver)
    else:
      pass # do nothing or log the error
  finally:
    release(driver)

And the calling code

workwithdevice(devicename, dowork)

Answer 2

This functionality seems to have been rejected. Python developers often prefer the explicit variant:

if need_more_workers():
    newcomm = get_new_comm(comm)
    # ...

You can also use higher-order functions:

def filter_comm(comm, nworkitems, callback):
    if foo:
        callback(get_new_comm())

# ...

some_local_var = 5
def do_work_with_newcomm(newcomm):
    # we can access the local scope here

filter_comm(comm, nworkitems, do_work_with_newcomm)

Answer 3

How about something like this instead:

@filter_comm(comm, nworkitems)
def _(newcomm):  # Name is unimportant - we'll never reference this by name.
    ... do work with communicator newcomm...

You implement the filter_comm decorator to do whatever work it should with comm and nworkitems, then based on those results decide whether to execute the function it's wrapped around or not, passing in newcomm.

It's not quite as elegant as with, but I think this is a bit more readable and closer to what you want than the other proposals. You could name the inner function something other than _ if you don't like that name, but I went with it since it's the normal name used in Python when the grammar requires a name which you'll never actually use.