List comprehension in Swift

Question

The language guide has revealed no trace of list comprehension. What\'s the neatest way of accomplishing this in Swift? I\'m looking for something similar t

Answer 1

Here's an extension to the Array types that combines filter and map into one method:

extension Array {

    func filterMap(_ closure: (Element) -> Element?) -> [Element] {

        var newArray: [Element] = []
        for item in self {
            if let result = closure(item) {
                newArray.append(result)
            }
        }
        return newArray
    }

}

It's similar to map except you can return nil to indicate that you don't want the element to be added to the new array. For example:

let items = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

let newItems = items.filterMap { item in
    if item < 5 {
        return item * 2
    }
    return nil
}

This could also be written more concisely as follows:

let newItems = items.filterMap { $0 < 5 ? $0 * 2 : nil }

In both of these examples, if the element is less than 5, then it is multiplied by two and added to the new array. If the closure returns nil, then the element is not added. Therefore, newIems is [2, 4, 6, 8].

Here's the Python equivalent:

items = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
newItems = [n * 2 for n in items if n < 5]

Answer 2

Got to admit, I am surprised nobody mentioned flatmap, since I think it's the closest thing Swift has to list (or set or dict) comprehension.

var evens = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9].flatMap({num -> Int? in 
    if num % 2 == 0 {return num} else {return nil}
})

Flatmap takes a closure, and you can either return individual values (in which case it will return an array with all of the non-nil values and discard the nils) or return array segments (in which case it will catenate all of your segments together and return that.)

Flatmap seems mostly (always?) to be unable to infer return values. Certainly, in this case it can't, so I specify it as -> Int? so that I can return nils, and thus discard the odd elements.

You can nest flatmaps if you like. And I find them much more intuitive (although obviously also a bit more limited) than the combination of map and filter. For example, the top answer's 'evens squared', using flatmap, becomes,

var esquares = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9].flatMap({num -> Int? in 
    if num % 2 == 0 {return num * num} else {return nil}
})

The syntax isn't quite as one-line not-quite-the-same-as-everything-else as python's is. I'm not sure if I like that less (because for the simple cases in python it's very short and still very readable) or more (because complex cases can get wildly out of control, and experienced python programmers often think that they're perfectly readable and maintainable when a beginner at the same company can take half an hour to puzzle out what it was intended to do, let alone what it's actually doing.)

Here is the version of flatMap from which you return single items or nil, and here is the version from which you return segments.

It's probably also worth looking over both array.map and array.forEach, because both of them are also quite handy.

Answer 3

As of Swift 2.x, there are a few short equivalents to your Python-style list comprehension.

The most straightforward adaptations of Python's formula (which reads something like "apply a transform to a sequence subject to a filter") involve chaining the map and filter methods available to all SequenceTypes, and starting from a Range:

// Python: [ x for x in range(10) if x % 2 == 0 ]
let evens = (0..<10).filter { $0 % 2 == 0 }

// Another example, since the first with 'x for x' doesn't
// use the full ability of a list comprehension:
// Python: [ x*x for x in range(10) if x % 2 == 0 ]
let evenSquared = (0..<10).filter({ $0 % 2 == 0 }).map({ $0 * $0 })

Note that a Range is abstract — it doesn't actually create the whole list of values you ask it for, just a construct that lazily supplies them on demand. (In this sense it's more like Python's xrange.) However, the filter call returns an Array, so you lose the "lazy" aspect there. If you want to keep the collection lazy all the way through, just say so:

// Python: [ x for x in range(10) if x % 2 == 0 ]
let evens = (0..<10).lazy.filter { $0 % 2 == 0 }
// Python: [ x*x for x in range(10) if x % 2 == 0 ]
let evenSquared = (0..<10).lazy.filter({ $0 % 2 == 0 }).map({ $0 * $0 })

Unlike the list comprehension syntax in Python (and similar constructs in some other languages), these operations in Swift follow the same syntax as other operations. That is, it's the same style of syntax to construct, filter, and operate on a range of numbers as it is to filter and operate on an array of objects — you don't have to use function/method syntax for one kind of work and list comprehension syntax for another.

And you can pass other functions in to the filter and map calls, and chain in other handy transforms like sort and reduce:

// func isAwesome(person: Person) -> Bool
// let people: [Person]
let names = people.filter(isAwesome).sort(<).map({ $0.name })

let sum = (0..<10).reduce(0, combine: +)

Depending on what you're going for, though, there may be more concise ways to say what you mean. For example, if you specifically want a list of even integers, you can use stride:

let evenStride = 0.stride(to: 10, by: 2) // or stride(through:by:), to include 10

Like with ranges, this gets you a generator, so you'll want to make an Array from it or iterate through it to see all the values:

let evensArray = Array(evenStride) // [0, 2, 4, 6, 8]

^{Edit: Heavily revised for Swift 2.x. See the edit history if you want Swift 1.x.}

Answer 4

Generally, a list comprehension in Python can be written in the form:

[f(x) for x in xs if g(x)]

Which is the same as

map(f, filter(g, xs))

Therefore, in Swift you can write it as

listComprehension<Y>(xs: [X], f: X -> Y, g: X -> Bool) = map(filter(xs, g), f)

For example:

map(filter(0..<10, { $0 % 2 == 0 }), { $0 })

Answer 5

As of Swift 2 you can do something like this:

var evens = [Int]()
for x in 1..<10 where x % 2 == 0 {
    evens.append(x)
}

// or directly filtering Range due to default implementations in protocols (now a method)
let evens = (0..<10).filter{ $0 % 2 == 0 }

Answer 6

One way would be :

var evens: Int[]()
for x in 0..<10 {
    if x%2 == 0 {evens += x} // or evens.append(x)
}

• Range operators
• Arrays