You don't have to handle all the cases

Oct 18, 2023

Scala has many nice features that make it a very expressive and powerful language. But whenever I go back to JS, I often find myself missing one small, often overlooked feature: partial functions.

Note: The target audience for this article are people just starting out with Scala.

If you ever dabbled your fingers in functional programming you probably learned that your functions must be pure, total and referential transparency all the way!

At first glance, the term itself may sound like an anti-pattern; why would we want our functions to be partial, in other words not defined for every argument they accept?

Pattern matching primer

Before we dive into partial functions, let’s take a look at how regular pattern matching works in Scala.

enum Pizza: // nominal type definition with three known instances, could be also defined with sealed trait and case classes in a Scala 2 fashion
  case Margherita(ham: Boolean)
  case Pepperoni(olives: Boolean)
  case Capricciosa(olives: Boolean, mozzarella: Boolean)

def hasOlives(pizza: Pizza) = pizza match 
  case Pizza.Pepperoni(olives)      => olives
  case Pizza.Capricciosa(olives, _) => olives
  case Pizza.Margherita(_)          => false

Above we have defined a Pizza ADT which might be useful in modeling the simplified domain a pizza ordering. The restaurant serves only three types of pizza, in each of them we can enable or disable particular ingredients. The chef was against changing anything else 🍕.

When the order is received, the kitchen needs to know what ingredients to expect, so we have defined a hasOlives function which takes our Pizza instance and returns a Boolean.

The match statement inside the function tests the Pizza instance against three cases. In the first case we check if the type of our Pizza is Pizza.Pepperoni and if so we extract the first constructor argument, the olives: Boolean and return it. Similarly in the second case, but here we discard the second constructor argument, mozzarella: Boolean. In the third case, we know that Margherita doesn’t have any olives, so we just return false, discarding the ham: Boolean.

Now we can imagine that somewhere further on there is a next step that might look like this:

def putOlivesOnPizza(order: Pizza): PizzaWithOlives = ???
// `???`` means no implementation, it compiles, but will throw an Error in runtime. Useful for prototyping.

val orders = List(
    Pizza.Pepperoni(olives = true),
    Pizza.Margherita(ham = true)
)

val pizzasWithOlives: List[PizzaWithOlives] = orders
  // execute `hasOlives` function for each collection element, returns a new collection with only the elements for which `hasOlives` returned `true`
  .filter(hasOlives)
  // for each element in the new collection execute the `putOlivesOnPizza` function
  .map(putOlivesOnPizza)

`collect` all the things!

For most situations this is fine, but whenever we see the filter- map pattern we can use the collect:

val pizzasWithOlives = orders.collect: // in Scala 3, single-argument functions can also be invoked by `:` and a newline
    case order: Pizza if hasOlives(order) => putOlivesOnPizza(order)

collect accepts something that looks like the pattern matching expression from before, but it’s also a bit different, so let’s unpack it:

First of all, with case order: Pizza we match on type without extracting any variables. This will essentially match all Pizza instances in our ADT.
Secondly, there is a pattern guard in the form of if <boolean expression>. It will further narrow the matched value to only those for which the expression returns the true.
We don’t handle the cases where hasOlives returns false, nor there is any fallback (like case _ => ???)

While the first two are not specific to partial functions, the third characteristic is a little bit suspicious. In regular pattern matching expressions we always want to handle all cases to preserve the exhaustiveness, but here - we don’t have to.

Turns out that instead of accepting normal functions like Function1[A, B] (unsugered A => B), the collect method accepts a PartialFunction[A, B]!

Essentially, the PartialFunction can state for which arguments it is defined using the case keyword

The collect defined above combines the filter and map into a single operation, without the need for two loops. The predicate that passed to filter is now taken from the left side of => in the partial function case, and the transformations in map are taken from the right side.

(partial) function composition

One of the essentials of functional programming is composition, in the general case it can be boiled down to the following:

val a: Int => Int = n => n + 1
val b: Int => Int = _ * 2 // same as `n => n * 2` but a little shorter

val c = a andThen b // or a.andThen(b), a mathy `b ∘ a`.
val c2: Function1[Int, Int] = n => b(a(n)) // We can write `Int => Int` or `Function1[Int, Int]`, the `1` in the name stands for number of arguments

println(c(2)) // 6
println(c2(2)) // 6

The first thing to notice in the above example is that we used val instead of def to define functions. In fact, the definitions using def are methods, and are slightly different from the actual functions, but for now this is of no consequence to us for now, so let’s ignore it and move on to the star of this snippet - the andThen method.

andThen is a method on a function object that accepts another function and a returns a third function that combines the previous two. It is equivalent to the c2 function definition.

We’ve used normal functions here, but it’s important to remember that the PartialFunction definition extends Function1, so anything the normal function can do, the partial function can do too.

However, there are some functionalities that are only possible with partial functions. Let’s have a look at the following, slightly larger example:

type Coupon = Int // a type alias, similar to the TypeScript one
type CouponHandler = PartialFunction[Coupon, Pizza]

// stub logic below
val priceWithCoupons: Pizza => BigDecimal     =  _ => BigDecimal(1) // the `_` indicates that we don't care about the argument
val priceWithoutCoupons: Coupon => BigDecimal =  _ => BigDecimal(2)

val handleXmasCoupons: CouponHandler =
  case 123 => Pizza.Margherita(ham = true)
  case 222 => Pizza.Margherita(ham = false) 
val handleLotteryCoupons: CouponHandler = 
  case 2213 | 2222 => Pizza.Pepperoni(olives = true) // the `|` is a pattern match alternative, the case will be matched on either `2213` or `2222`. Don't confuse it with union types.
  case 233         => Pizza.Pepperoni(olives = false)
val handleTwitchStreamCoupons: CouponHandler = ???
val handleMealDealCoupons: CouponHandler     = ???

val finalPrice: Coupon => BigDecimal = handleXmasCoupons
    .orElse(handleLotteryCoupons)
    .orElse(handleTwitchStreamCoupons)
    .orElse(handleMealDealCoupons)
    .andThen(priceWithCoupons)
    .applyOrElse(_, priceWithoutCoupons) // the `_` will refer to the function argument, `Coupon`, this is the same syntax magic like in the previous snipped

println(finalPrice(233)) // 1
println(finalPrice(17))  // 2

In this code fragment we can see our first, stand alone partial function definition, the CouponHandler. We will create a few instances of it, and then bind them all together with the .orElse method.

orElse allows us to provide a fallback partial function to be called if the argument passed is not in the domain of the first function. This in turn allows us to break a large pattern match into smaller, more digestible bits, which can then be used to construct a single, larger partial function. Nice and tidy.

The andThen, described already above, allows as to add a post-processingstep that is common to all cases. Here it takes a Pizza produced by partial function and returns its price.

At the end of our processing pipeline we call applyOrElse. This is another partial function specific method. It allows us to call the function with an argument, but if the argument can’t be handled then the fallback function priceWithoutCoupons is called.

Things to look out for

Let’s define a partial function without all the syntax niceties.

val fraction = new PartialFunction[Double, Double]: // an anonymous instance of the `PartialFunction[Double, Double]`
 def apply(n: Double): Double = 1 / n
 def isDefinedAt(n: Double): Boolean = n != 0

println(fraction.isDefinedAt(3)) // true
println(fraction(3)) // 0.3333333333333333

println(fraction.isDefinedAt(0)) // false
println(fraction(0)) // Infinity, ups...

To have a partial function instance, we need to implement two methods:

apply, similar to the regular function, it will host the actual implementation.
isDefinedAt, which tells us which arguments can be handled by our function.

We can immediately see that there is a discrepancy between the apply' and isDefinedAt’ methods. We can call the former even if the isDefinedAt returns false for the same argument. This is perhaps the biggest problem with partial functions.

Consider the following example:

def foo(fn: Double => Double) = fn(0)

println(foo(fraction)) // Infinity, again...

Since a partial function is also a normal function (it extends Function1), you can pass it to anywhere you want. This can be problematic in situations like the one above. The way you code a normal function and the way you code a partial function is inherently different. In the latter, you deliberately don’t handle all cases and instead rely on the isDefinedAt or case guards.

We can mitigate the problem by using the lift method.

println(fraction.lift(3)) // Some(0.3333333333333333)
println(fraction.lift(0)) // None

lift allows us to call a function in the same way as you would do with apply / (), but if the argument is not in the domain of our partial function we will get None, otherwise it will be a Some(result).

However, this doesn’t solve the underlying problem, which is the interface-expectation mismatch.

Summary

I believe that pattern matching and partial functions in Scala are wonderful tools that unlock powerful composition patterns and can make your code cleaner and easier to work with. However, as mentioned above, passing partial functions all over the codebase can be error-prone and I’d rather limit them to smaller,-ish, enclosed scopes.