Python函数中变量返回值数量不确定的优雅赋值方式

Question

我正在使用一个返回结果为列表的Python库函数。根据实验配置，该列表由一个或两个元素组成。在这种情况下，将函数输出分配给变量的一种简单方法如下：

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)

然而，这种方式看起来有些啰嗦，因为我们需要先将bbox和segm都初始化为None，然后再判断 if len(results) > 1。有没有一种更Pythonic的方法可以避免这种情况呢？理想情况下，以下代码会非常好:

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

- gab

Answer 1

您可以在Python 3.10中使用结构化模式匹配：

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

然而，如果test_module总是返回一个元组，你可以使用拆包：

bbox, *segm = test_model()

现在，如果text_model返回一个值，则segm将是一个空列表（[]）。但是，如果它返回两个值，则segm将包含一个元素。因此，您的完整代码可以如下：

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

- Ajax1234

---

非常好的答案！非常感谢！我想我会选择拆包，因为我正在使用Python < 3.10。同时使用'match'似乎是一个非常好的、灵活的解决方案 :) 我猜方法重载变得更加简洁了，对吧？ - gab

Answer 2

如果您使用的是Python 3.10，请选择Ajax1234的答案。

否则，您的首要任务是编写易于阅读的内容，而不是进行任何花哨的操作。像这样将值赋给元组即可：

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

这种方式有很多优点：

• 一目了然地看到正在发生的情况。
• 如果列表长度不为1或2，你将得到一个错误，而不是它悄无声息地失败。

在处理异常时，使用异常来处理实际上是异常的情况是有意义的。因此，如果您预计一个列表项会有更多的结果，您可以反过来写：

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

- Joooeey

Answer 3

你可以扩展return list的一个额外元素，然后切掉需要的两个元素。

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

Pythonic的方式是像@Ajax1234提到的那样使用*解包，但这是一种替代方法。

如果您控制函数定义，也可以添加一个装饰器。

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

- Axe319

Answer 4

两种更常见的方法是，如果test_model()只返回一个值，则将segm赋值为None。

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

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

- Kelly Bundy