Wrapping a generic method in a class extension

Question

I'm trying to wrap the generic Array method compactMap inside an Array extension to give the purpose of the method more meaning/readability. I'm simply trying to take an Array of Optionals and remove any and all nil values from it.

extension Array {
    public func removeNilElements() -> [Element] {
        let noNils = self.compactMap { $0 }
        return noNils // nil values still exist
    }
}

The problem I am having is that compactMap here is not working. nil values are still in the resulting Array noNils. When I use the compactMap method directly without using this wrapper, I get the desired result of an Array with no nil values.

let buttons = [actionMenuButton, createButton]   // [UIBarButtonItem?]
let nonNilButtons = buttons.compactMap { $0 }    // works correctly
let nonNilButtons2 = buttons.removeNilElements() // not working

Am I not designing my extension method correctly?

Answer 1

You have to define the method for an array of optional elements, and the return type as the corresponding array of non-optionals. This can be done with a generic function:

extension Array {
    public func removingNilElements<T>() -> [T] where Element == T?    {
        let noNils = self.compactMap { $0 }
        return noNils
    }
}

Example:

let a = [1, 2, nil, 3, nil, 4]   // The type of a is [Int?]
let b = a.removingNilElements()  // The type of b is [Int]
print(b) // [1, 2, 3, 4]

In your code, $0 has the (non-optional) type Element, and it just wrapped by the compiler into an optional in order to match the argument type of compactMap().