List comprehension in Swift

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.


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 })

With Swift 5, you can choose one of the seven following Playground sample codes in order to solve your problem.

#1. Using stride(from:to:by:) function

let sequence = stride(from: 0, to: 10, by: 2)
let evens = Array(sequence)
// let evens = sequence.map({ $0 }) // also works
print(evens) // prints [0, 2, 4, 6, 8]

#2. Using Range filter(_:) method

let range = 0 ..< 10
let evens = range.filter({ $0 % 2 == 0 })
print(evens) // prints [0, 2, 4, 6, 8]

#3. Using Range compactMap(_:) method

let range = 0 ..< 10
let evens = range.compactMap({ $0 % 2 == 0 ? $0 : nil })
print(evens) // prints [0, 2, 4, 6, 8]

#4. Using sequence(first:next:) function

let unfoldSequence = sequence(first: 0, next: {
    $0 + 2 < 10 ? $0 + 2 : nil
})
let evens = Array(unfoldSequence)
// let evens = unfoldSequence.map({ $0 }) // also works
print(evens) // prints [0, 2, 4, 6, 8]

#5. Using AnySequence init(_:) initializer

let anySequence = AnySequence<Int>({ () -> AnyIterator<Int> in
    var value = 0
    return AnyIterator<Int> {
        defer { value += 2 }
        return value < 10 ? value : nil
    }
})
let evens = Array(anySequence)
// let evens = anySequence.map({ $0 }) // also works
print(evens) // prints [0, 2, 4, 6, 8]

#6. Using for loop with where clause

var evens = [Int]()
for value in 0 ..< 10 where value % 2 == 0 {
    evens.append(value)
}
print(evens) // prints [0, 2, 4, 6, 8]

#7. Using for loop with if condition

var evens = [Int]()
for value in 0 ..< 10 {
    if value % 2 == 0 {
        evens.append(value)
    }
}
print(evens) // prints [0, 2, 4, 6, 8]

Tags:

List Comprehension

Swift

Recent Posts