让Python函数根据参数数量干净地返回标量或列表

Question

免责声明：我在寻找一个适用于Python 2.6的解决方案，如果有的话。

我想要一个函数，这个函数在接收一个值时返回一个单一的结果，而在接收多个值时返回一系列的结果：

>>> a = foo(1)
2
>>> b, c = foo(2, 5)
>>> b
3
>>> c
6

为了更清楚，这样做是为了让一些函数调用看起来更美观，而不是：

a, = foo(1)

或者

a = foo(1)[0]

现在，比较笨拙的解决方案大概是这样的：

def foo(*args):
    results = [a + 1 for a in args]
    return results if len(results) > 1 else results[0]

有没有什么语法糖（或者函数）可以让这个看起来更简洁？像下面这样的？

def foo(*args):
    return *[a + 1 for a in args]

Answer 1

这个代码可以处理0个或多个参数，我想这正是你想要的。

def foo(*args):
    return map(lambda x: x + 1, args) or [None]

补充说明：我修改了一下，增加了一个None列表，以防传入0个参数时出现问题。

Answer 2

如果你觉得这样更好，可以写一个装饰器来简化那个if语句：

import functools
def unpacked(method):
    @functools.wraps(method)
    def _decorator(*args):
        result = method(*args)
        return results if len(results) != 1 else results[0]
    return _decorator

使用方法：

@unpacked
def foo(*args):
    return [arg + 1 for arg in args]

Answer 3

你可以很简单地写一个叫做 scalify 的函数，如果列表里只有一个元素，它就返回这个元素。这个名字的意思就是把它变成一个标量（单一值）。

def scalify(l):
    return l if len(l) > 1 else l[0]

然后你可以像这样在你的函数里使用它：

def foo(*args):
    return scalify([a + 1 for a in args])

这样做是可以的，但我和那些建议你不要这样做的人是同一派的。原因之一是，这样做会让你无法遍历结果，除非你确定传入了至少两个项目。此外，如果你有一个列表，在调用这个函数时，你必须把列表拆开，这样就失去了它作为“列表”的特性，而且你也不能保证会得到一个列表回来。这些缺点在我看来远远超过了这种方法可能带来的任何好处。