Swift中处理可为nil、元素不足n的数组的前n个元素获取方法

刚好之前处理过类似需求，用Swift的语言特性就能写出非常优雅的实现，完全不用冗余的条件判断。

核心实现：函数式写法 #

最直接的方式是写一个工具函数，结合可选链和标准库的prefix(_:)方法：

func takeFirst<T>(_ count: Int, from optionalArray: [T]?) -> [T]? {
    optionalArray?.prefix(count).map(Array.init)
}

为什么这行代码能满足所有需求？ #

• 处理nil输入：optionalArray? 开启可选链，如果原数组是nil，整个表达式直接返回nil，完美匹配你要的nil输入返回nil的规则；
• 元素不足n的情况：prefix(count) 是Swift标准库方法，当数组元素数量小于传入的count时，会直接返回整个数组，比如数组[1]传入count=2，就返回[1]；
• 转换为数组：prefix返回的是ArraySlice<T>，用map(Array.init)把它转换成真正的[T]类型，避免后续使用切片可能带来的意外。

测试用例验证 #

// 正常数组取前2个
let fullArray = [1,2,3,4,5]
takeFirst(2, from: fullArray) // 输出: Optional([1, 2])

// 数组元素不足count
let shortArray = [1]
takeFirst(2, from: shortArray) // 输出: Optional([1])

// 数组为nil
let nilArray: [Int]? = nil
takeFirst(2, from: nilArray) // 输出: nil

更Swift的风格：扩展Optional #

如果想让调用更自然，可以给Optional写个扩展，让所有包装了集合类型的可选值都能用这个方法：

extension Optional where Wrapped: Collection {
    func takeFirst(_ count: Int) -> [Wrapped.Element]? {
        self?.prefix(count).map(Array.init)
    }
}

调用起来就像原生方法一样流畅：

fullArray.takeFirst(2) // Optional([1, 2])
shortArray.takeFirst(2) // Optional([1])
nilArray.takeFirst(2) // nil

这个扩展还能复用在其他集合类型上，比如字符串取前几个字符：

let optionalString: String? = "Hello World"
optionalString.takeFirst(5) // Optional(["H", "e", "l", "l", "o"])

内容的提问来源于stack exchange，提问作者gmoraleda