CODEWITHSUNDEEP
pythonassign
Gabriele
Gabriele

Reputation: 949

Elegant way to assign a variable number of returns from a python function

I'm using a python library with a function returning results as a list. Depending on the experiment configuration, this list is made of one or two elements. In this case, a straightforward way to assign the function output is the following:

bbox, segm = None, None
results = test_model()  # returns a list with either one or two elements
bbox = results[0]
if len(results) > 1:
    segm = results[1]

# [some other code here]
if segm is not None:
    plot(segm)

However, this seems to be quite verbose because we need to first initialise both bbox and segm to None, and then evaluate if len(results) > 1. Is there any pythonic way to avoid this? Ideally, something like this would be very nice:

bbox, segm = test_model()  # returns a list with either one or two elements
if segm is not None:
    plot(segm)

Upvotes: 3

Views: 430

Answers (4)

Kelly Bundy
Kelly Bundy

Reputation: 27609

Two more ways, assigning None to segm if test_model() returns only one value:

bbox, segm, *_ = *test_model(), None

bbox, segm = (*test_model(), None)[:2]

Upvotes: 0

Joooeey
Joooeey

Reputation: 3876

If you're on Python 3.10, go with Ajax1234's answer.

Otherwise, your priority is to write something that's easy to read instead of doing anything fancy. Assigning to a tuple like this would work:

try:
    bbox, segment = results
except ValueError:
    bbox, segment = results[0], None

This has a bunch of advantages:

  • You see at a glance what is happening.
  • If the list has any other length than 1 or 2, you'll get an error instead of it failing silently.

When working with exceptions it makes sense to use the exception for the case that is actually exceptional. So if you expect way more results with one list item, you could write it the other way around:

try:
    bbox, = results
    segment = None
except ValueError:
    bbox, segment = results

Upvotes: 2

Axe319
Axe319

Reputation: 4365

You could extend the return list an additional element and then slice off the 2 needed elements.

def foo(boolean):
    if boolean:
        return [True, True]
    return [False]

bar, baz = (foo(False) + [None])[:2]
print(bar, baz) # False None

bar, baz = (foo(True) + [None])[:2]
print(bar, baz) # True True

The pythonic way would be * unpacking as @Ajax1234 mentioned, but this is an alternative.

If you have control of the function definition you could add a decorator as well.

def returns_two(func):
    def wrapper(*args):
        return (func(*args) + [None])[:2]
    return wrapper


@returns_two
def foo(boolean):
    if boolean:
        return [True, True]
    return [False]


bar, baz = foo(False)
print(bar, baz) # False None

bar, baz = foo(True)
print(bar, baz) # True True

Upvotes: 1

Ajax1234
Ajax1234

Reputation: 71451

You can use structural pattern matching, available in Python 3.10:

match test_model():
   case [bbox, segm]:
      plot(segm)
   case bbox:
      pass

However, if test_module always returns a tuple, you can use unpacking:

bbox, *segm = test_model()

Now, if text_model returned one value, segm will be an empty list ([]). However, if it returned two values, segm will contain a single element. Thus, your full code can become:

bbox, *segm = test_model()
if segm:
   plot(segm[0])

Upvotes: 4

Related Questions