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:
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