Scala: Filtering based on type

Just wanted to expand on mkneissl's answer with a "more generic" version that should work on many different collections in the library:

scala> import collection._
import collection._

scala> import generic.CanBuildFrom
import generic.CanBuildFrom

scala> def partition[X,A,B,CC[X] <: Traversable[X], To, To2](xs : CC[X])(f : X => Either[A,B])(
     |   implicit cbf1 : CanBuildFrom[CC[X],A,To], cbf2 : CanBuildFrom[CC[X],B,To2]) : (To, To2) = {
     |   val left = cbf1()
     |   val right = cbf2()
     |   xs.foreach(f(_).fold(left +=, right +=))
     |   (left.result(), right.result())
     | }
partition: [X,A,B,CC[X] <: Traversable[X],To,To2](xs: CC[X])(f: (X) => Either[A,B])(implicit cbf1: scala.collection.generic.CanBuildFrom[CC[X],A,To],implicit cbf2: scala.collection.generic.CanBuildFrom[CC[X],B,To2])(To, To2)

scala> partition(List(1,"two", 3)) {                                                                
     |   case i: Int => Left(i)                                                                     
     |   case x => Right(x)                                                                         
     | }
res5: (List[Int], List[Any]) = (List(1, 3),List(two))

scala> partition(Vector(1,"two", 3)) {
     |   case i: Int => Left(i)       
     |   case x => Right(x)           
     | }
res6: (scala.collection.immutable.Vector[Int], scala.collection.immutable.Vector[Any]) = (Vector(1, 3),Vector(two))

Just one note: The partition method is similar, but we need to capture a few types:

X -> The original type for items in the collection.

A -> The type of items in the left partition

B -> The type of items in the right partition

CC -> The "specific" type of the collection (Vector, List, Seq etc.) This must be higher-kinded. We could probably work around some type-inference issues (see Adrian's response here: http://suereth.blogspot.com/2010/06/preserving-types-and-differing-subclass.html ), but I was feeling lazy ;)

To -> The complete type of collection on the left hand side

To2 -> The complete type of the collection on the right hand side

Finally, the funny "CanBuildFrom" implicit paramters are what allow us to construct specific types, like List or Vector, generically. They are built into to all the core library collections.

Ironically, the entire reason for the CanBuildFrom magic is to handle BitSets correctly. Because I require CC to be higher kinded, we get this fun error message when using partition:

scala> partition(BitSet(1,2, 3)) {    
     |   case i if i % 2 == 0  => Left(i)
     |   case i if i % 2 == 1 => Right("ODD")
     | }
<console>:11: error: type mismatch;
 found   : scala.collection.BitSet
 required: ?CC[ ?X ]
Note that implicit conversions are not applicable because they are ambiguous:
 both method any2ArrowAssoc in object Predef of type [A](x: A)ArrowAssoc[A]
 and method any2Ensuring in object Predef of type [A](x: A)Ensuring[A]
 are possible conversion functions from scala.collection.BitSet to ?CC[ ?X ]
       partition(BitSet(1,2, 3)) {

I'm leaving this open for someone to fix if needed! I'll see if I can give you a solution that works with BitSet after some more play.


EDITED

While using plain partition is possible, it loses the type information retained by collect in the question.

One could define a variant of the partition method that accepts a function returning a value of one of two types using Either:

import collection.mutable.ListBuffer

def partition[X,A,B](xs: List[X])(f: X=>Either[A,B]): (List[A],List[B]) = {
  val as = new ListBuffer[A]
  val bs = new ListBuffer[B]
  for (x <- xs) {
    f(x) match {
      case Left(a) => as += a
      case Right(b) => bs += b
    }
  }
  (as.toList, bs.toList)
}

Then the types are retained:

scala> partition(List(1,"two", 3)) {
  case i: Int => Left(i)
  case x => Right(x)
}

res5: (List[Int], List[Any]) = (List(1, 3),List(two))

Of course the solution could be improved using builders and all the improved collection stuff :) .

For completeness my old answer using plain partition:

val (a,b) = x partition { _.isInstanceOf[SomeClass] }

For example:

scala> val x = List(1,2, "three")
x: List[Any] = List(1, 2, three)

scala> val (a,b) = x partition { _.isInstanceOf[Int] }
a: List[Any] = List(1, 2)
b: List[Any] = List(three)

Use list.partition:

scala> val l = List(1, 2, 3)
l: List[Int] = List(1, 2, 3)

scala> val (even, odd) = l partition { _ % 2 == 0 }
even: List[Int] = List(2)
odd: List[Int] = List(1, 3)

EDIT

For partitioning by type, use this method:

def partitionByType[X, A <: X](list: List[X], typ: Class[A]): 
    Pair[List[A], List[X]] = {
    val as = new ListBuffer[A]
    val notAs = new ListBuffer[X]
    list foreach {x =>
      if (typ.isAssignableFrom(x.asInstanceOf[AnyRef].getClass)) {
        as += typ cast x 
      } else {
        notAs += x
      }
    }
    (as.toList, notAs.toList)
}

Usage:

scala> val (a, b) = partitionByType(List(1, 2, "three"), classOf[java.lang.Integer])
a: List[java.lang.Integer] = List(1, 2)
b: List[Any] = List(three)