Python：確定作為函數參數傳遞的列表中的元素數

Question

我有一些功能foo(list)

我想

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

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

在不知道列表本身是什么的情況下如何做到這一點，理想情況下，僅使用assert語句該如何做？

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