Grouping elements of a Sequence in Swift
Lately, I’ve been looking for ways to share some of the things I’ve learned in a different, more entertaining way. I’m pretty happy with the results so far. It really made me hungry for experimenting with drawing, painting, recording videos, podcasting. The problem is, I’m still pretty bad at it and that’s why I won’t show you the result of my last 30 minutes spent in Pixelmator.
Oops.
Today I had to work on a relatively common task: grouping a list of view models into separate sections and display them in a UITableView
. Somehow that kind of task always ends up with all the other one that doesn’t need to be made generic or reusable easily. After all, it never takes more than 5 to 10 lines of code. Let’s make it both generic and reusable anyway.
What do we want
For this post, I’m going to reuse a piece of code I’m using in Synchronizable’s tests (more about that project later). We have a list of heroes and we want to group them by brand:
enum Brand: String {
case Marvel
case DC
case Valiant
}
struct Hero {
let name: String
let brand: Brand
}
let heroes: [Hero] = [
Hero(name: "Spider-Man", brand: .Marvel),
Hero(name: "Batman", brand: .DC),
Hero(name: "Zephyr", brand: .Valiant),
Hero(name: "Green-Lantern", brand: .DC),
Hero(name: "Dr Strange", brand: .Marvel),
Hero(name: "Blue Beetle", brand: .DC)
]
In a world where we wouldn’t mind solving this problem over and over again, we would do it like this:
var groupedHeroes = [Brand: [Hero]]()
heroes.forEach { hero in
if groupedHeroes[hero.brand] == nil {
groupedHeroes[hero.brand] = []
}
groupedHeroes[hero.brand]?.append(hero)
}
Pretty straightforward. Note that I’m not saying something like that should be made available as a Pod or an iOS framework. That’s way too simple. It doesn’t hurt to have this kind of snippet around, though.
About Hashable
Let’s step back a little and talk about Dictionaries in Swift. A Dictionary is an association between a key and a value. In our situation, the key is the brand and the value is a collection of heroes. You can store pretty much anything in a dictionary, the only restriction is having a key that conforms to the Hashable protocol as we need unique keys. As the documentation states:
Many types in the standard library conform to Hashable: strings, integers, floating-point and Boolean values, and even sets provide a hash value by default. Your own custom types can be
Hashable
as well.
Making it generic
The Swift standard library has a bunch of methods we can draw inspiration from to make our own groupBy
method. The index(where:)
method available in the Collection
type is a great one: it takes a closure that will allow us to filter a sequence and then retrieve the index. We’ll use a similar approach to retrieve the key
we want the sequence to be grouped by.
We want our groupBy
method to leverage Swift’s type system so we don’t need to cast the resulting dictionary. Fortunately, that’s exactly what generics constraints are for. That way we can reuse the code we wrote earlier, make it generic and available as an extension of the Sequence type which Array inherits from!
extension Sequence {
// Using a `typealias` because it's shorter to write `E`
// Think of it as a shortcut
typealias E = Iterator.Element
// Declaring a `K` generic that we'll use as the type of the key
// for the resulting dictionary. The only restriction is having
// it conforming to the `Hashable` protocol
func groupBy<K: Hashable>(handler: (E) -> K) -> [K: [E]] {
// Creating the resulting dictionary
var grouped = [K: [E]]()
// Iterating over our elements
self.forEach { item in
// Retrieving the key based on the current item
let key = handler(item)
if grouped[key] == nil {
grouped[key] = []
}
grouped[key]?.append(item)
}
return grouped
}
}
So yeah. If generics are not solving your problems, make sure you’re using enough generics.