How can I do 'if..else' inside a for-comprehension?

Question

I am asking a very basic question which confused me recently. I want to write a Scala For expression to do something like the following:

for (i <- expr1) {
  if (i.method) {
    for (j <- i) {
      if (j.method) {
        doSomething()
      } else {
        doSomethingElseA()
      }
    }
  } else {
    doSomethingElseB()
  }
}

The problem is that, in the multiple generators For expression, I don't know where can I put each for expression body.

for {i <- expr1
  if(i.method) // where can I write the else logic ?
  j <- i 
  if (j.method)
} doSomething()

How can I rewrite the code in Scala Style?

Solution 1

The first code you wrote is perfectly valid, so there's no need to rewrite it. Elsewhere you said you wanted to know how to do it Scala-style. There isn't really a "Scala-style", but I'll assume a more functional style and tack that.

for (i <- expr1) {
  if (i.method) {
    for (j <- i) {
      if (j.method) {
        doSomething()
      } else {
        doSomethingElseA()
      }
    }
  } else {
    doSomethingElseB()
  }
}

The first concern is that this returns no value. All it does is side effects, which are to be avoided as well. So the first change would be like this:

val result = for (i <- expr1) yield {
  if (i.method) {
    for (j <- i) yield {
      if (j.method) {
        returnSomething()
        // etc

Now, there's a big difference between

for (i <- expr1; j <- i) yield ...

and

for (i <- expr1) yield for (j <- i) yield ...

They return different things, and there are times you want the later, not the former. I'll assume you want the former, though. Now, before we proceed, let's fix the code. It is ugly, difficult to follow and uninformative. Let's refactor it by extracting methods.

def resultOrB(j) = if (j.method) returnSomething else returnSomethingElseB
def nonCResults(i) = for (j <- i) yield resultOrB(j)
def resultOrC(i) = if (i.method) nonCResults(i) else returnSomethingC
val result = for (i <- expr1) yield resultOrC(i)

It is already much cleaner, but it isn't returning quite what we expect. Let's look at the difference:

trait Element
object Unrecognized extends Element
case class Letter(c: Char) extends Element
case class Punct(c: Char) extends Element
val expr1 = "This is a silly example." split "\\b"

def wordOrPunct(j: Char) = if (j.isLetter) Letter(j.toLower) else Punct(j)
def validElements(i: String) = for (j <- i) yield wordOrPunct(j)
def classifyElements(i: String) = if (i.nonEmpty) validElements(i) else Unrecognized
val result = for (i <- expr1) yield classifyElements(i)

The type of result there is Array[AnyRef], while using multiple generators would yield Array[Element]. The easy part of the fix is this:

val result = for {
  i <- expr1
  element <- classifyElements(i)
} yield element

But that alone won't work, because classifyElements itself returns AnyRef, and we want it returning a collection. Now, validElements return a collection, so that is not a problem. We only need to fix the else part. Since validElements is returning an IndexedSeq, let's return that on the else part as well. The final result is:

trait Element
object Unrecognized extends Element
case class Letter(c: Char) extends Element
case class Punct(c: Char) extends Element
val expr1 = "This is a silly example." split "\\b"

def wordOrPunct(j: Char) = if (j.isLetter) Letter(j.toLower) else Punct(j)
def validElements(i: String) = for (j <- i) yield wordOrPunct(j)
def classifyElements(i: String) = if (i.nonEmpty) validElements(i) else IndexedSeq(Unrecognized)
val result = for {
  i <- expr1
  element <- classifyElements(i)
} yield element

That does exactly the same combination of loops and conditions as you presented, but it is much more readable and easy to change.

About Yield

I think it is important to note one thing about the problem presented. Let's simplify it:

for (i <- expr1) {
  for (j <- i) {
    doSomething
  }
}

Now, that is implemented with foreach (see here, or other similar questions and answer). That means the code above does exactly the same thing as this code:

for {
  i <- expr1
  j <- i
} doSomething

Exactly the same thing. That is not true at all when one is using yield. The following expressions do not yield the same result:

for (i <- expr1) yield for (j <- i) yield j

for (i <- expr1; j <- i) yield j

The first snippet will be implemented through two map calls, while the second snippet will use one flatMap and one map.

So, it is only in the context of yield that it even makes any sense to worry about nesting for loops or using multiple generators. And, in fact, generators stands for the fact that something is being generated, which is only true of true for-comprehensions (the ones yielding something).

Solution 2

The part

for (j <- i) {
   if (j.method) {
     doSomething(j)
   } else {
     doSomethingElse(j)
   }
 }

can be rewritten as

for(j <- i; e = Either.cond(j.method, j, j)) {
  e.fold(doSomething _, doSomethingElse _)  
}  

(of course you can use a yield instead if your do.. methods return something)

Here it is not so terrible useful, but if you have a deeper nested structure, it could...

Solution 3

import scalaz._; import Scalaz._

val lhs = (_ : List[X]) collect { case j if j.methodJ => doSomething(j) } 
val rhs = (_ : List[X]) map doSomethingElse
lhs <-: (expr1 partition methodI) :-> rhs

Solution 4

You can not. The for(expr; if) construct just filter the element that must be handled in the loop.

Solution 5

If the order isn't important for the calls to doSomething() and doSomethingElse() then you can rearrange the code like this.

val (tmp, no1) = expr1.partition(_.method)
val (yes, no2) = tmp.partition(_.method)

yes.foreach(doSomething())
no1.foreach(doSomethingElse())
no2.foreach(doSomethingElse())

To answer your original question, I think that for comprehensions can be quite nice for specific use cases, and your example doesn't fit nicely.

View more solutions