Chaining PartialFunctions with andThen in Scala

Question

Let us reuse examples from Daily scala :

type PF = PartialFunction[Int,Int]

val pf1 : PF = {case 1 => 2}                      

val pf2 : PF = {case 2 => 3}                      

and let us add:

val pf3 : PF = {case 3 => 4}

andThen works as expected here:

pf1 andThen pf2 isDefinedAt(x)

returns true iff x == 1 (actually, pf2 does not need to be a PartialFunction at all)

However, I expected that:

pf1 andThen pf3 isDefinedAt(x)

would return false for all x (i.e., iff pf1 is defined, check for pf3), but it does not and only validates pf1.

In the end, pf1 andThen pf3 lift(x) always result in a MatchError. I would prefer to get None… I can obtain this behavior by lifting each function such as in pf1.lift(x).flatMap(pf3.lift) but is there any easier way using pure PartialFunction API? (and without lifting each partial function individually?)

Answer 1

If you look at andThen:

def andThen[C](k: (B) => C): PartialFunction[A, C]

This composes the receiver with a function and not a partial function. That is, k is expected to be fully defined, it doesn't have isDefinedAt. Therefore, the resulting partial function does not need to alter the behaviour of isDefinedAt, it will still just has to consult the first partial function.

You could write your own extension that composes two partial functions:

implicit class ComposePartial[A, B](pf: PartialFunction[A, B]) {
  def collect[C](that: PartialFunction[B, C]): PartialFunction[A, C] =
    new PartialFunction[A, C] {
      def apply(a: A): C = that(pf(a))
      def isDefinedAt(a: A) = pf.isDefinedAt(a) && {
        val b = pf(a)
        that.isDefinedAt(b)
      }
    }
}

pf1 collect pf2 isDefinedAt(1)  // true
pf1 collect pf3 isDefinedAt(1)  // false

The problem is that you have to invoke pf(a), so given that Scala doesn't enforce purity, you may end up executing side effects unwantedly.

Answer 2

You need the equivalent of flatMap for PartialFunctions.

implicit class CollectPartial[A, B](f: PartialFunction[A, B]) {
    def collect[C](g: PartialFunction[B, C]) = Function.unlift { a: A =>
        f.lift(a).flatMap(g.lift)
    }
}

Use it like

val a: PartialFunction[String, Int] = ...
val b: PartialFunction[Int, Char] = ...
val c: PartialFunction[String, Char] = a collect b

This works as expected even with side-effects.