Playing around with Kotlin Sealed Classes

While working on my Master's thesis, I recently had to configure probabilities. These values were organized in a tree-like structure. Because of this, I thought of using inheritance - and with Kotlin already in place, using sealed classes seemed at least like an interesting opportunity to me.

Disclaimer: This post is just for fun and for showing some cool Kotlin-specific things. In the end, using Jackson to deserialize input into data classes is probably more effective.

Sample

Let's say, we want to configure the likeliness in percent of some actions using YAML for instance. This could look like the following:

Configuration:
    Greeting:
        SayHello: 70
        WaveAt: 80
        Hug: 5
    Talking:
        DoSmallTalk: 90
        Insult: 1

As you can see, there are two subsections. However, this can easily get an arbitrary depth.

Before we'll reproduce that structure, we should clarify what sealed classes are.

Sealed Classes

The official documentation calls them extensions to enum classes, used for representing restricted class hierarchies.

Sealed classes are declared using the keyword sealed, are abstract and must have only private constructors.

Simply put, a sealed class ensures, that all possible subtypes are known at compile time. This comes especially handy, when used in conjunction with Kotlin's when clause:

sealed class TruthOrDare
class Truth(val question: String) : TruthOrDare()
class Dare(val task: String) : TruthOrDare()

fun nextTurn(input: TruthOrDare) = when(input) {
    is Truth -> println("Answer the following question: ${input.question}")
    is Dare -> println("Your task is: ${input.task}")
    // the `else` clause is not required because we've covered all the cases
}

Sealed Classes in Action

Let's come back to the original task: using sealed classes for parsing configuration.

The structure of the above YAML snippet can be represented as follows:

sealed class Configuration {

    sealed class Greeting : Configuration() {

        object SayHello : Greeting()
        object WaveAt : Greeting()
        object Hug : Greeting()

    }

    sealed class Talking : Configuration() {

        object DoSmallTalk : Talking()
        object Insult : Talking()

    }

}

We have nested sealed subclasses and objects (a shortcut for Singletons in Kotlin). We'll see later, why we are using objects here.

However, yet there is nothing configured. So how do we use this structure?

The answer is: Reflection.

Digression

Before we come to the actual structure though, I would like to introduce a small wrapper class, to encapsulate our findings:

class Probabilities private constructor(
        private val backingMap: MutableMap<Configuration, Int> = mutableMapOf()
) : Map<Configuration, Int> by backingMap {

    constructor(configuration: Probabilities.() -> Unit) : this() {
        configuration()
    }

    infix fun Configuration.withProbabilityOf(percent: Int) = backingMap.put(this, percent)

    override fun toString(): String = backingMap.entries.joinToString { "${it.key::class.simpleName} = ${it.value}" }
}

In this class you can see some other cool Kotlin features: implementation by delegation, receiver functions and infix functions.

Implementing the Map interface can be a pain in Java. Kotlin provides a smart solution with implementation by delegation. All you need to do, is to use the by keyword in conjunction with a value, that already implements the interface. So in this case we simply delegate all Map-specific operations to a MutableMap, which is created by the default constructor.

To make configuration of the Probabilities class easy, we provide a second constructor taking a receiver function. This means, that the this keyword of the provided lambda points to an instance of Probabilities class.

Another receiver function is also provided by withProbabilityOf. However, this is also an infix function, which is marked by the infix keyword and enables you to write something like:

SayHello withProbabilityOf 99

Parsing

Finally, here is the code to parse e.g. a Map of String and Any:

fun fromMap(configuration: Map<String, Any>) = Probabilities { // (1)
    @Suppress("UNCHECKED_CAST")
    fun parseFor( // (2)
            configuration: Map<String, Any>, // (3)
            parent: KClass<out Configuration>?,
            clazz: KClass<out Configuration>
    ) {
        // (4)
        if (configuration.containsKey(clazz.simpleName) && parent?.isSuperclassOf(clazz) != false) {
            when (val value = configuration[clazz.simpleName]) {
                is Map<*, *> -> { // (5)
                    clazz.sealedSubclasses.forEach { subclass ->
                        parseFor(value as Map<String, Any>, clazz, subclass)
                    }
                }
                is Int -> { // (6)
                    if (clazz.objectInstance == null)
                        throw RuntimeException("${clazz.simpleName} should be an object")

                    clazz.objectInstance!! withProbabilityOf value
                }
                else -> throw RuntimeException("unknown property ${clazz.simpleName}")
            }
        }
    }

    parseFor(configuration, null, Configuration::class)
}

Some explanation (mind the comments in code):

• (1) we construct a new instance of Probabilities by making use of one of Kotlin's greatest features: if the last parameter of any function or constructor is a lambda function, we can omit parentheses and only use curly braces
• (2) another beautiful thing are local functions: since we don't need the recursive function parseFor anywhere else, we simply declare it inside our fromMap function
• (3) the parameters of our parseFor function are:
  • the subtree of the configuration
  • the assumed parent as specified by the configuration
  • the current class as specified by the configuration's property
• (4) if somewhere on top level of the current subtree the expected class is found and the assumed parent matches the real superclass, then we use when on the property value's type
• (5) if it's a Map, apply parseFor on each sealed subclass
• (6) if it's an Int and clazz is an object, we configure the probability to be that value

In the above code, you might see, why having leaves being objects is a good idea: there will always be exactly one instance. So they are perfect for being used as a key in a map.

Sweet. - But why?

To be honest, all of the above might seem like an over-engineered construct. Compared to the example using sealed classes, an equivalent structure with data classes would be similar and could be easily used with e.g. Jackson out of the box:

data class Configuration(
    val greeting: Greeting?,
    val talking: Talking?
) {

    data class Greeting(
        val sayHello: Int?,
        val waveAt: Int?,
        val hug: Int?
    )

    data class Talking(
        val doSmallTalk: Int?,
        val insult: Int?
    )
}

This is the great thing about Kotlin: you have so many language constructs that help you build your software the best way, whatever this means. Therefore, it's necessary to play around with and get to know these possibilities.

Using Jackson

To make use of the concept of using sealed classes in conjunction with Jackson, we need to provide a custom deserializer. Since we already have the mapping logic however, this is an easy task:

class ProbabilitiesDeserializer : StdDeserializer<Probabilities>(Probabilities::class.java) {

    override fun deserialize(p: JsonParser?, ctxt: DeserializationContext?) =
            p?.readValueAs<Map<String, Any>>(object : TypeReference<Map<String, Any>>() {})
                    ?.let { fromMap(it) }

}

For this and the ability to parse YAML we need the following dependencies:

• com.fasterxml.jackson.module:jackson-module-kotlin:2.11.2
• com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.11.2

Finally, we are ready to create an ObjectMapper and read our YAML configuration:

val mapper = ObjectMapper(YAMLFactory()).apply {
    registerModule(SimpleModule().addDeserializer(Probabilities::class.java, ProbabilitiesDeserializer()))
    registerModule(KotlinModule())
}
val probabilities = mapper.readValue<Probabilities>(yamlInput)
println(probabilities.toString())
println("Say hello with probability of ${probabilities[SayHello]} percent.")

The result will be:

SayHello = 70, WaveAt = 80, Hug = 5, DoSmallTalk = 90, Insult = 1
Say hello with probability of 70 percent.

That's it.

Bonus: Ensuring all properties are set

One last thing I wanted to add, is how you enforce all objects to be configured.

All we have to do is to find all the leaves of our configuration tree and compare them to the input's content:

val leaves = leavesOf(Configuration::class)
fun leavesOf(baseClass: KClass<out Configuration>): List<KClass<out Configuration>> =
    if (!baseClass.isSealed) {
        listOf(baseClass)
    } else {
        baseClass.sealedSubclasses.flatMap(::leavesOf)
    }

fun Probabilities.ensureAllActionsCovered() {
    val keys = keys.map { it::class }
    val unconfigured = leaves.filter { leaf -> !keys.contains(leaf) }
    if (unconfigured.isNotEmpty())
        throw RuntimeException("Unconfigured leaves: ${unconfigured.joinToString()}")
}

Finally, we've seen a last cool Kotlin feature: extension functions. They allow us to add some functionality to an otherwise closed class. Like a bonus in a way.

You call them as if they were a class method:

probabilities.ensureAllActionsCovered()

Closing notes

Thanks for reading, I hope you liked it. You can find all the code in a kscript on Github as a Gist.