How to simplify this kind of Python code?

Question

I have many blocks of code like this:

try:
  a = get_a()
  try:
    b = get_b()

    # varying codes here, where a and b are used

  finally:
    if b:
      cleanup(b)
finally:
  if a:
    cleanup(a)

I hope to write some magic code like this:

some_magic:
  # varying codes here, where a and b are also available

Is this possible?

Answer 1

If you can't or don't want to implement context protocol for a and b, you can use contextlib facilities to make a context:

from contextlib import contextmanager

@contextmanager
def managed(a):
    try:
        yield a
    finally:
        if a:
            cleanup(a)

with managed(get_a()) as a, managed(get_b()) as b:
    # do something here
    pass

Answer 2

Let classes of a and b implement the context manager protocol, and use the with statement:

with get_a() as a, get_b() as b:
    do_magic

Now if get_a and get_b returned open file handles, they'd be automatically closed at the end of the block. If the values returned are of a custom class, that class should have the __enter__ and __exit__ magic methods.