How to Write a Scala Class That Can Be Used in a ‘for’ Expression

“You must unlearn what you have learned.”

Yoda, Star Wars


As a reminder, the reason for the next several lessons on for expressions is so that you can get to a point where you can understand the following code:

def stackManip: State[Stack, Int] = for {
    _ <- push(3)
    a <- pop
    b <- pop
} yield(b)

To make this code easier to understand — or understandable at all, depending on your background — let’s start digging into how the Scala for expression works.

How for expressions work

The book, Programming in Scala, co-written by Martin Odersky, is the definitive reference for the Scala programming language, and Section 23.6 of that book, “Generalizing for,” describes the translation rules for how the Scala compiler converts (a) the for expressions you write into (b) a series of method calls that may include map, flatMap, foreach, and withFilter.

Just as for expressions are compiled into these four functions, the opposite is also true: if you write a class that implements these functions, it can be used inside a for expression.

Rules about how these functions enable for

In the next several lessons I’ll show how to write your own custom data types that can be used in for expressions. To do this you need to know the rules that govern how a for expression is enabled by these functions. Programming in Scala gives us these translation rules:

  1. If a custom data type defines a foreach method, it allows for loops (both with single and multiple generators). (Note the emphasis on the word “loops” in that definition. This refers to the simple Java-style use of a for loop, i.e., for (i <- ints) println(i).)
  2. If a data type defines only map, it can be used in for expressions consisting of a single generator.
  3. If it defines flatMap as well as map, it allows for expressions consisting of multiple generators.
  4. If it defines withFilter, it allows for filter expressions starting with an if within the for expression.

While a for expression technically doesn’t require specific signatures for each of these methods, if you want your custom class to work well with a for expression it should generally implement those methods with the signatures shown in this example class:

abstract class CustomClass[A] {
    def map[B](f: A => B): CustomClass[B]
    def flatMap[B](f: A => CustomClass[B]): CustomClass[B]
    def withFilter(p: A => Boolean): CustomClass[A]
    def foreach(b: A => Unit): Unit

If those type signatures don’t make sense just yet, fear not, I’ll cover them in the next few lessons.

And now — because we’re at a point where we really need to know how for expressions work — in the next few lessons I’ll demonstrate how to satisfy these rules as we build a custom class that can be used inside a Scala for expression.

books i’ve written