Generate class diagrams with a Kotlin DSL for PlantUML
Tags: kotlin, documentation, plantuml, dsl
Published: 2021-12-18 15:05:27

Technical documentation of the systems helps us and our (future) colleagues to understand how things are built and how they work. The importance of internal documentation was recently also proven by the [State of DevOps 2021](https://cloud.google.com/blog/products/devops-sre/announcing-dora-2021-accelerate-state-of-devops-report), which notes it as one of its key findings this year. However, while probably most of us understand the relevance of good documentation, it is still not trivial to always keep it in sync with reality. Deadlines are tight, people want value delivered asap, "we will update it next sprint", and more decisions to get short term gains.

In a recent project, we therefore wanted to automate the creation of certain parts of the documentation. The first automation candidate we identified were class diagrams. As part of our documentation-as-code approach, these diagrams were already manually made with [PlantUML](https://plantuml.com/), an open-source tool allowing users to create diagrams from a plain text language. For example, with PlantUML, we can write the following text:

```plantuml
@startuml
class Person {
  + id : Int
  + name : String
  + getName() : String
  + getPets() : List<Pet>
}

class Pet {
  + id : Int
  + name : String
  + owner : Person
  + getName() : String
  + getOwner() : Person
}

Person - Pet : owns >
@enduml
```

The PlantUML tool can take this simple text document and generate a PNG with a cool retro-ish look:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/gswfdnwtog2utg1diu23.png)
 
With PlantUML, you can update the documentation in the same pull request as your code changes and benefit from git's version-control. Also, as its part of the pull requests, colleagues can actually review the diagrams and so each change to documentation is approved by someone else. I personally think this is already a big improvement compared to stashing your documentation somewhere far away from your code on a Confluence page or, even worse, a Sharepoint.

For the project that I worked on, we would have these PlantUML diagrams in the repository and a `git push` would trigger a pipeline that creates the PNG images and deploys everything to our [Sphinx](https://www.sphinx-doc.org/en/master/)-based company documentation server. That works well, but we still needed to manually update these diagrams and that felt like applying our changes twice. Once to the code and once to the documentation. Could we not just generate the docs?

# The plan
We want to create class diagrams automatically. A first thought would be to use StringBuilder. While that is perfectly possible, we preferably had a more structured way that also ensures a checked and standardised output. To this end, we explored creating a Kotlin DSL (Domain-Specific Language) for PlantUML diagrams. This article explains our approach and the decisions we took. Below you see a sneak-peak of the syntax. Do you like what you see? Then please read on!

```kotlin
fun main() {
    val myDiagram = classDiagram {
        clazz("Person") {
            public (
                Field("id", "Integer"),
           		Method("getPets", "List<Pet>")
            )
        }

        clazz("Pet") {
            public (
                Field("id", "Integer"),
                Method("getName", "String"),
                Method("getOwner", "Person")
            )

            private {
            	Field("name", "String")
            }
        }

        relationship("Person", "Pet", RelationshipType.ASSOCIATION, "is owned by")
	}
}
```

# Lambda with receivers
The secret ingredient to create DSL's in Kotlin is its receiver types. Using these, a lambda with a receiver allows you to call methods of an object in the body of a lambda without any qualifiers. Much has already been written about receivers in other articles so, in this article, we will only look at the applying lambda's with receivers to our specific use case.

# Diagram-as-a-constrained-tree
The key idea behind our PlantUML DSL is to see a class diagram as a very constrained tree structure. Each node has children of a very specific type. For example, a class can only have fields and methods.

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/q9fzmaz2tagi85gvu69y.png)

DSL's are great to encode such constraints. Let's start by defining an abstract base class for all our elements. `PlantUmlElement` has a list to keep track of its children which are all subtypes of `PlantUmlElement` as well. A new child can be added through the `addChild` method, which has two implementations:

1. The first implementation accepts just the new child as parameter. This method will be used if that child is a leaf node and does not require further initialisation.
2. Otherwise, we also pass a lambda with receiver to initialise the grandchildren. 

We also require subclasses to implement the `toString()` method to get the PlantUML String. I admit that this is a slight misuse of `toString()`, but it does make our syntax cleaner for our example.

```kotlin
abstract class PlantUmlElement {
    protected val children = mutableListOf<PlantUmlElement>()

    protected fun <T: PlantUmlElement> addChild(child: T) {
        children.add(child)
    }

    protected fun <T: PlantUmlElement> addChild(child: T, init: T.() -> Unit) {
        child.init()
        this.addChild(child)
    }

    abstract override fun toString(): String
}
```

# Implementation of the elements
All elements essentially follow the same approach:
 
1. Extend the `PlantUmlElement` base class.
2. If the element is not a leaf, then expose a method for each possible child element type that calls `addChild`. If the child element type is also not a leaf, then the variant of `addChild` is called with a lambda with receiver. 
3. Override `toString()` to return the PlantUML representation of this element. If the element is not a leaf, then call the `toString()` methods of the children as well.

The implementation then looks like this:

```kotlin
enum class AttributeVisibility { PUBLIC, PRIVATE, PROTECTED }

abstract class Attribute(private val first: String, private val second: String? = null): PlantUmlElement() {
    private var visibility: AttributeVisibility = AttributeVisibility.PUBLIC

    fun setVisibility(visibility: AttributeVisibility): Attribute {
        this.visibility = visibility
        return this
    }

    protected fun getModifier() =
        when(this.visibility) {
            AttributeVisibility.PUBLIC -> "+"
            AttributeVisibility.PRIVATE -> "-"
            AttributeVisibility.PROTECTED -> "*"
        }

    abstract override fun toString(): String
}

class Field(private val name: String, private val dataType: String? = null): Attribute(name, dataType) {
    override fun toString() = "${getModifier()} $name ${if(dataType != null) ": $dataType" else ""}"
}

class Method(private val name: String, private val returnType: String? = null): Attribute(name, returnType) {
    override fun toString() = "${getModifier()} $name() ${if(returnType != null) ": $returnType" else ""}"
}

enum class RelationshipType { EXTENSION, COMPOSITION, AGGREGATION, ASSOCIATION }

class Relationship(private val left: String, private val right: String, private val type: RelationshipType, private val label: String? = null): PlantUmlElement() {
    override fun toString() = "$left ${getArrow()} $right ${if (label != null) ": $label" else ""}"

    private fun getArrow() =
        when(this.type) {
            RelationshipType.EXTENSION -> "<|--"
            RelationshipType.COMPOSITION -> "*--"
            RelationshipType.AGGREGATION -> "o--"
            RelationshipType.ASSOCIATION -> "<--"
        }
}

class Class(private val name: String, private val identifier: String): PlantUmlElement() {
    private fun addAttributes(attributes: List<Attribute>, visibility: AttributeVisibility)
        = attributes.forEach { addChild(it.setVisibility(visibility)) }

    fun public(vararg attributes: Attribute) = addAttributes(attributes.toList(), AttributeVisibility.PUBLIC)
    fun private(vararg attributes: Attribute) = addAttributes(attributes.toList(), AttributeVisibility.PRIVATE)

    override fun toString() = "class \"${this.name}\" as ${this.identifier} {\n\t${children.joinToString("\n\t")}\n}"
}

class ClassDiagram: PlantUmlElement() {
    fun clazz(name: String, identifier: String = name.replace(" ", ""), init: Class.() -> Unit) =
        addChild(Class(name, identifier), init)

    fun relationship(left: String, right: String, type: RelationshipType, label: String? = null) =
        addChild(Relationship(left, right, type, label))

    override fun toString() = "@startuml ldm\n\t${children.joinToString("\n\n")}\n@enduml"
}
```

All that we now still require is a function to start with the root `ClassDiagram`. This function is not part of a class and accepts a lambda with receiver that is then applied to a freshly created `ClassDiagram` object. 

```kotlin
fun classDiagram(init: ClassDiagram.() -> Unit) = ClassDiagram().apply(init)
```

# A more complex example
Using the DSL we just created, we can make class diagrams that are slightly more complex than what we already saw in this article. For example, we can write the following diagram:

```kotlin
fun main() {
    val myDiagram = classDiagram {
        clazz("Person") {
            public (
                Field("id", "Integer"),
                Method("getName", "String")
            )
        }

        clazz("Customer") {
            public (
                Field("id", "Integer"),
                Field("emailAddress", "String"),
                Method("getName", "String"),
                Method("getAccountBalance", "Integer"),
                Method("verifyPassword", "Boolean")
            )

            private (
                Field("name", "String"),
                Field("password", "String")
            )
        }

        clazz("Wallet") {
            public (
                Field("id", "Integer"),
                Field("customerId", "Integer"),
                Method("getBalance", "Integer")
            )
        }

        clazz("Order") {
            public (
                Field("id", "Integer"),
                Field("customerId", "Integer"),
                Method("place", "Integer"),
                Method("getTotalPrice", "Float")
            )
        }

        clazz("OrderItem") {
            public (
                Field("id", "Integer"),
                Field("orderId", "Integer"),
                Method("getDescription", "String")
            )

            private (
                Field("description", "String"),
                Field("price", "Float")
            )
        }

        relationship("Person", "Customer", RelationshipType.EXTENSION, "extends")
        relationship("Customer", "Wallet", RelationshipType.AGGREGATION, "belongs to")

        relationship("Order", "OrderItem", RelationshipType.COMPOSITION, "is part of")

        relationship("Order", "Customer", RelationshipType.ASSOCIATION, "places")
    }

    println(myDiagram)
}
```

Which results in this image:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/9lc10p4htw78lox71fwo.png)

# Et Voila!
That's all we need to create a simple class diagram. Obviously, PlantUML has a lot more features and these can be integrated into our DSL as well. To actually automate documentation generation, we would then use `forEach`, `map` and other language features to dynamically put such diagram together. At least, we now have an approach that is more robust that just building a string using `StringBuilder` or just appending Strings. We can enforce constraints and the generation code is also much more readable.

The full code and all examples in this article can be found on GitHub: [BjornvdLaan/plantuml-dsl-kotlin](https://github.com/BjornvdLaan/plantuml-dsl-kotlin).