Kotlin annotation processor example

annotation_build.gradle

plugins {
    id 'org.jetbrains.kotlin.jvm' version '1.3.61'
}

version 'unspecified'

repositories {
    mavenCentral()
}

dependencies {
    implementation "org.jetbrains.kotlin:kotlin-stdlib-jdk8"
}

compileKotlin {
    kotlinOptions.jvmTarget = "1.8"
}
compileTestKotlin {
    kotlinOptions.jvmTarget = "1.8"
}

codegen_build.gradle

plugins {
    id 'org.jetbrains.kotlin.jvm' version '1.3.61'
    id 'org.jetbrains.kotlin.kapt' version '1.3.61'
}

version 'unspecified'

repositories {
    mavenCentral()
}

dependencies {
    implementation project(':annotation')
    implementation "org.jetbrains.kotlin:kotlin-stdlib-jdk8"
    implementation 'com.squareup:kotlinpoet:1.5.0'
    implementation "com.google.auto.service:auto-service:1.0-rc6"
    kapt "com.google.auto.service:auto-service:1.0-rc6"
}

compileKotlin {
    kotlinOptions.jvmTarget = "1.8"
}
compileTestKotlin {
    kotlinOptions.jvmTarget = "1.8"
}

FeatureFlag.kt

package com.example.annotation.annotation

@Retention(AnnotationRetention.SOURCE)
@Target(AnnotationTarget.FUNCTION)
@MustBeDocumented
annotation class FeatureFlag(val key: String, val defaultValue: Boolean)

FeatureFlagCodegen.kt

package com.example.annotation.codegen

import com.example.annotation.annotation.FeatureFlag
import com.example.annotation.annotation.FeatureFlagGroup
import com.google.auto.service.AutoService
import com.squareup.kotlinpoet.*
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import java.io.File
import javax.annotation.Generated
import javax.annotation.processing.AbstractProcessor
import javax.annotation.processing.Processor
import javax.annotation.processing.RoundEnvironment
import javax.annotation.processing.SupportedOptions
import javax.lang.model.SourceVersion
import javax.lang.model.element.*
import javax.tools.Diagnostic

private const val KAPT_KOTLIN_GENERATED_OPTION_NAME = "kapt.kotlin.generated"

@AutoService(Processor::class)
@SupportedOptions(KAPT_KOTLIN_GENERATED_OPTION_NAME)
class FeatureFlagCodegen : AbstractProcessor() {

    override fun getSupportedAnnotationTypes(): Set<String> = setOf(
        FeatureFlagGroup::class.java.name,
        FeatureFlag::class.java.name
    )

    override fun getSupportedSourceVersion(): SourceVersion = SourceVersion.latestSupported()

    override fun process(annotations: MutableSet<out TypeElement>, roundEnv: RoundEnvironment): Boolean {
        if (roundEnv.processingOver()) return false
        roundEnv.getElementsAnnotatedWith(FeatureFlagGroup::class.java)
            .filter { it.kind == ElementKind.INTERFACE }
            .forEach { featureFlagGroupElement ->
                val featureFlagElements = featureFlagGroupElement.enclosedElements
                    .filter { it.getAnnotation(FeatureFlag::class.java) != null && it.kind == ElementKind.METHOD }
                val packageName = processingEnv.elementUtils.getPackageOf(featureFlagGroupElement).toString()
                generateImpl(packageName, featureFlagGroupElement, featureFlagElements)
            }
        return roundEnv.getElementsAnnotatedWith(FeatureFlagGroup::class.java).any { it.kind == ElementKind.INTERFACE }
    }

    private fun generateImpl(
        packageName: String,
        featureFlagGroupElement: Element,
        featureFlagElements: List<Element>
    ) {
        val generatedSourcesRoot = processingEnv.options[KAPT_KOTLIN_GENERATED_OPTION_NAME].orEmpty()
        if (generatedSourcesRoot.isEmpty()) {
            processingEnv.messager.printMessage(
                Diagnostic.Kind.ERROR,
                "Can't find the target directory for generated Kotlin files."
            )
            return
        }

        val implClassName = "${featureFlagGroupElement.simpleName}Impl"
        val file = File(generatedSourcesRoot)
        file.mkdir()

        val className = ClassName(
            (featureFlagGroupElement.enclosingElement as PackageElement).qualifiedName.toString(),
            featureFlagGroupElement.simpleName.toString()
        )

        val funSpecs = featureFlagElements.map {
            val executableElement = it as ExecutableElement
            val featureFlagAnnotation = it.getAnnotation(FeatureFlag::class.java)
            FunSpec.overriding(executableElement)
                .addStatement("return ${featureFlagAnnotation.defaultValue}")
                .build()
        }

        val defaultValuesMapCodeBlock =
            featureFlagElements.foldIndexed(CodeBlock.builder().add("mapOf(\n")) { index, builder, element ->
                val featureFlagAnnotation = element.getAnnotation(FeatureFlag::class.java)
                builder.add("%S·to·${featureFlagAnnotation.defaultValue}", featureFlagAnnotation.key)
                if (index < featureFlagElements.size - 1) {
                    builder.add(",\n")
                } else {
                    builder.add("\n)")
                }
                builder
            }.build()

        FileSpec.builder(packageName, implClassName)
            .addType(
                TypeSpec.classBuilder(implClassName)
                    .addAnnotation(
                        AnnotationSpec.builder(Generated::class.java)
                            .addMember("value = [%S]", FeatureFlagCodegen::class.java.name)
                            .build()
                    )
                    .addKdoc(CodeBlock.of("Concrete implementation of [%T].", className))
                    .primaryConstructor(
                        FunSpec.constructorBuilder()
                            .addAnnotation(ClassName("javax.inject", "Inject"))
                            .build()
                    )
                    .addSuperinterface(className)
                    .addFunctions(funSpecs)
                    .addType(
                        TypeSpec.companionObjectBuilder(null)
                            .addProperty(
                                PropertySpec.builder(
                                    "defaultValues",
                                    Map::class.asClassName()
                                        .parameterizedBy(String::class.asClassName(), Boolean::class.asClassName())
                                )
                                    .initializer(defaultValuesMapCodeBlock)
                                    .addKdoc("Default value map")
                                    .build()
                            )
                            .build()
                    )
                    .build()
            )
            .build()
            .writeTo(file)
    }
}

FeatureFlagGroup.kt

package com.example.annotation.annotation

@Retention(AnnotationRetention.SOURCE)
@Target(AnnotationTarget.CLASS)
@MustBeDocumented
annotation class FeatureFlagGroup

Main.kt

package com.example.annotation.sample

import com.example.annotation.sample.example.MyFeatureFlagsA
import com.example.annotation.sample.example.MyFeatureFlagsAImpl

fun main() {
    val featureFlagsA: MyFeatureFlagsA = MyFeatureFlagsAImpl()
    val featureFlagsB: MyFeatureFlagsB = MyFeatureFlagsBImpl()

    println("featureFlagsA.featureA1 = ${featureFlagsA.featureA1()}")
    println("featureFlagsA.featureA2 = ${featureFlagsA.featureA2()}")
    println("featureFlagsA.featureA3 = ${featureFlagsA.featureA3()}")
    println("featureFlagsA.featureA4 = ${featureFlagsA.featureA4()}")
    println("featureFlagsB.featureB1 = ${featureFlagsB.featureB1()}")
    println("featureFlagsB.featureB2 = ${featureFlagsB.featureB2()}")
}

MyFeatureFlagsA.kt

package com.example.annotation.sample

import com.example.annotation.annotation.FeatureFlag
import com.example.annotation.annotation.FeatureFlagGroup

@FeatureFlagGroup
interface MyFeatureFlagsA {

    @FeatureFlag(key = "feature_a1", defaultValue = true)
    fun featureA1(): Boolean

    @FeatureFlag(key = "feature_a2", defaultValue = false)
    fun featureA2(): Boolean

    @FeatureFlag(key = "feature_a3", defaultValue = true)
    fun featureA3(): Boolean

    @FeatureFlag(key = "feature_a4", defaultValue = false)
    fun featureA4(): Boolean
}

MyFeatureFlagsB.kt

package com.example.annotation.sample

import com.example.annotation.annotation.FeatureFlag
import com.example.annotation.annotation.FeatureFlagGroup

@FeatureFlagGroup
interface MyFeatureFlagsB {
    @FeatureFlag(key = "feature_b1", defaultValue = false)
    fun featureB1(): Boolean

    @FeatureFlag(key = "feature_b2", defaultValue = false)
    fun featureB2(): Boolean
}

sample_build.gradle

plugins {
    id 'org.jetbrains.kotlin.jvm' version '1.3.61'
    id 'org.jetbrains.kotlin.kapt' version '1.3.61'
}

group 'com.example.annotation'
version '1.0-SNAPSHOT'

repositories {
    mavenCentral()
}

dependencies {
    implementation "org.jetbrains.kotlin:kotlin-stdlib-jdk8"
    implementation project(':annotation')
    kapt project(':codegen')
    implementation "javax.inject:javax.inject:1"
}

compileKotlin {
    kotlinOptions.jvmTarget = "1.8"
}
compileTestKotlin {
    kotlinOptions.jvmTarget = "1.8"
}