python - Python：确定作为函数参数传递的列表中的元素数量

Question

我有一些功能foo(list)

我想

assert foo(list with one element) == that one element

assert foo(list with 100 elements) == foo(list)

在不知道列表本身是什么的情况下，最好只使用断言语句，我该如何做到这一点？

Answer 1

if len(x) == 1:
    assert foo(x) == x[0]
elif len(x) == 100:
    assert foo(x) == x

或者，如果您希望确定foo的返回值始终具有上述属性，您可以这样做：

def foo(x):
    # compute return value `retval` based on `x`
    if len(x) == 1:
        assert retval == x[0]
    elif len(x) == 100:
        assert retval == x
    return retval