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 =&         


        

Answer 1

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.

Answer 2

Why not simply :

def compose[A,B,C](f: PartialFunction[A, B], g: PartialFunction[B, C]) : PartialFunction[A, C] =
Function.unlift(f.andThen(g.lift))

Answer 3

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.