Composing transformations with for... yield

Using the same db: Map[String, Int] phrase, containing the ages of different people, the following code is a naive implementation of a function that returns the average age of two people:

def averageAgeA(name1: String, name2: String, db: Map[String, Int]): Option[Double] = {
  val optOptAvg: Option[Option[Double]] =
    db.get(name1).map(age1 =>
      db.get(name2).map(age2 =>
        (age1 + age2).toDouble / 2))
  optOptAvg.flatten
}
val db = Map("John" -> 25, "Rob" -> 40)
averageAge("John", "Rob", db)
// res6: Option[Double] = Some(32.5)
averageAge("John", "Michael", db)
// res7: Option[Double] = None

The function returns Option[Double]. If name1 or name2 cannot be found in the db map, averageAge returns None. If both names are found, it returns Some(value). The implementation uses map to transform the value contained in the option. We end up with a nested Option[Option[Double]], but our function must return Option[Double]. Fortunately, we can use flatten to remove one level of nesting.

We managed to implement averageAge, but we can improve it using flatMap, as shown in the following code:

def averageAgeB(name1: String, name2: String, db: Map[String, Int]): Option[Double] =
  db.get(name1).flatMap(age1 =>
    db.get(name2).map(age2 =>
      (age1 + age2).toDouble / 2))

As its name suggests, flatMap is equivalent to composing flatten and map. In our function, we replaced map(...).flatten with flatMap(...).

So far, so good, but what if we want to get the average age of three or four people? We would have to nest several instances of flatMap, which would not be very pretty or readable. Fortunately, Scala provides a syntactic sugar that allows us to simplify our function further, called the for comprehension, as shown in the following code:

def averageAgeC(name1: String, name2: String, db: Map[String, Int]): Option[Double] =
  for {
    age1 <- db.get(name1)
    age2 <- db.get(name2)
  } yield (age1 + age2).toDouble / 2

When you compile a for comprehension, such as for { ... } yield { ... }, the Scala compiler transforms it into a composition of flatMap/map operations. Here is how it works:

• Inside the for block, there can be one or many expressions phrased as variable <- context, which is called a generator. The left side of the arrow is the name of a variable that is bound to the content of the context on the right of the arrow.
• Every generator except the last one is transformed into a flatMap expression.
• The last generator is transformed into a map expression.
• All context expressions (the right side of the arrow) must have the same context type.

In the preceding example, we used Option for the context type, but for yield can also be used with any class that has a flatMap and map operation. For instance, we can use for..yield with Vector to run nested loops, as shown in the following code:

for {
  i <- Vector("one", "two")
  j <- Vector(1, 2, 3)
} yield (i, j)
// res8: scala.collection.immutable.Vector[(String, Int)] = 
// Vector((one,1), (one,2), (one,3), (two,1), (two,2), (two,3))