简体   繁体   中英

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)

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])

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

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

还有两种方法,如果test_model()只返回一个值, segm None分配给segm

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

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM