DEV Community

Cover image for Generate class diagrams with a Kotlin DSL for PlantUML
BjornvdLaan
BjornvdLaan

Posted on • Updated on

Generate class diagrams with a Kotlin DSL for PlantUML

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, 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, an open-source tool allowing users to create diagrams from a plain text language. For example, with PlantUML, we can write the following text:

@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
Enter fullscreen mode Exit fullscreen mode

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

Image description

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-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!

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")
    }
}
Enter fullscreen mode Exit fullscreen mode

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

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.

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
}
Enter fullscreen mode Exit fullscreen mode

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:

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"
}
Enter fullscreen mode Exit fullscreen mode

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.

fun classDiagram(init: ClassDiagram.() -> Unit) = ClassDiagram().apply(init)
Enter fullscreen mode Exit fullscreen mode

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:

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)
}
Enter fullscreen mode Exit fullscreen mode

Which results in this image:

Image description

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.

Top comments (0)