Python context manager: conditionally executing body?

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.

Solution

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:
    handledrivererror(geterrordetails(driver))
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:
      handledrivererror(geterrordetails(driver))
  finally:
    release(driver)

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

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

# this works but you don't win much
@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:
    handledrivererror(geterrordetails(driver))

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.

Instead of not yielding we can raise an Exception ourselves:

# this works but you don't win much
@contextlib.contextmanager
def contextgetdriver(devicename):
  driver = getdriver(devicename)
  try:
    if isgooddriver(driver):
      yield driver
    else:
      raise NoGoodDriverException(geterrordetails(driver))
  finally:
    release(driver)

but now you need to handle the exception:

try:
  with contextgetdriver(devicename) as driver:
    dowork(driver)
except NoGoodDriverException as e:
  handledrivererror(e.errordetails)

which has practically the same cost of code complexity as the explicit testing for good driver.

Side note: with an exception we can decide to not handle it here and instead let it bubble up the call stack and handle it elsewhere. also there is a difference: by the time we handle the exception the driver has already been released. While with the explicit check the driver has not been released. (the except is outside of the with statement while the else is inside the with statement)

Now to the alternatives

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

# this works but please don't
def generatorgetdriver(devicename):
  driver = getdriver(devicename)
  try:
    if isgooddriver(driver):
      yield driver
    else:
      handledrivererror(geterrordetails(driver))
  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 that the only way to prevent the boilerplate code is with a callback

# this works fine but it is no longer a context manager really
def workwithdevice(devicename, callback):
  driver = getdriver(devicename)
  try:
    if isgooddriver(driver):
      callback(driver)
    else:
      handledrivererror(geterrordetails(driver))
  finally:
    release(driver)

And the calling code

workwithdevice(devicename, dowork)