Suma
Suma

Reputation: 34393

Work with collection keeping the type information

I have a following extension to Traversable:

import scala.reflect.ClassTag

object ContainerHelpers {

  implicit class TraversableOps[X](a: Traversable[X]) {
    def partitionByType[T <: X: ClassTag]: (Traversable[T], Traversable[X]) = {
      val (typed, other) = a.partition {
        case _: T => true
        case _ => false
      }
      (typed.map(_.asInstanceOf[T]), other)
    }

  }

  // test case:
  trait T
  case class A(i: Int) extends T

  val s = Seq[T](A(0), A(1))
  val sa = s.partitionByType[A]
  sa // I would like this to be Seq[A], not Traversable[A]

}

The extension works fine, however it does not keep a collection type - it always returns Traversable. I would like to write it so that partitionByType would use the type information known by partition and map, therefore type of sa should be (Seq[A], Seq[T]), not (Traversable[A], Traversable[T]). How can I do this, perhaps using higher-kinded types, or something like that?

Upvotes: 2

Views: 121

Answers (2)

mdm
mdm

Reputation: 3988

You can use CanBuildFrom and build the collections yourself, instead of relying on TraversableLike's implementation and then cast things down. This has the advantage of not doing any casting anywhere, and it does not cost you anything more than the alternative, because the partition method in TraversableLike does the same.

This is a first stab at getting what you want (I changed the parametrized types' names for clarity):

implicit class TraversableOps[Elem, T[_Ignore] <: Traversable[_Ignore]](a: T[Elem]) {
  def partitionByType[Target <: Elem: ClassTag]
     (implicit cbf1: CanBuildFrom[T[Elem], Elem, T[Elem]],
       cbf2: CanBuildFrom[T[Target], Target, T[Target]]) : (T[Target], T[Elem]) = {

    val l = cbf2()
    val r = cbf1()

    for (x <- a) x match{
      case t: Target => l += t
      case _ => r += x
    }

    (l.result, r.result)
 }
}

And it works with the following test case:

// test case:
trait A
trait B extends A
case class T(i: Int) extends A
case class U(i: Int) extends B

val s1 : List[A] = List(T(0), U(1))
val s2 : Seq[A] = List(T(0), U(1))

val sa1 = s1.partitionByType[T]  //sa1 : (List[T], List[A]) = (List(T(0)),List(U(1)))
val sa2 = s2.partitionByType[T]   //sa2 : (Seq[T], Seq[A]) = (List(T(0)),List(U(1)))

Upvotes: 3

chengpohi
chengpohi

Reputation: 14217

Use F-bounded container to do this, like:

  //F[_] extend from Traversable for own the original type of colleciton
  implicit class TraversableOps[X, F[_] <: Traversable[_]](a: F[X]) {
    def partitionByType[T <: X : ClassTag]: (F[T], F[X]) = {
      //cast to Traversable for known partition type 
      val (typed, other) = a.asInstanceOf[Traversable[X]].partition {
        case _: T => true
        case _ => false
      }
      //cast to the target type, Since the original is F-bounded, so when cast also is cast type safe
      (typed.asInstanceOf[F[T]], other.asInstanceOf[F[X]])
    }

  }

  // test case:
  trait T

  case class A(i: Int) extends T

  case class B(i: Double) extends T

  val s = Seq[T](A(0), A(1), B(2.0f))
  val sa: (Seq[A], Seq[T]) = s.partitionByType[A]
  println(sa._1)
  println(sa._2)

Upvotes: 2

Related Questions