Kotlinpoet による Kotlinコードの生成

LLM時代には、ルールベースで決まるコードはスクリプトで生成し、それを使ってLLMにコーディングしてもらうのが一つの効率的なやり方だと思っています。今回は、Kotlinpoetを使ったKotlinコード生成について勉強会で話したので共有させていただきます。

概要

https://square.github.io/kotlinpoet/

KotlinPoet is a Kotlin and Java API for generating .kt source files.

• できること: Kotlinのコード自体を生成する
  • Pythonのjinja2やQuarkusのquteなどのテンプレートエンジンを使うと、構文的に間違っている状態になりがち
• できないこと: ktファイルを読み込む
  • Typescriptのts-morphはASTを操作するので読み込むことができる
• 主な利用用途: データベースのスキーマなどのmetadataファイルからコードを生成する

例: https://square.github.io/kotlinpoet/#example

生成AIとの組み合わせ

生成AIと相性が良い

• 生成AIでKotlinpoetを使ったスクリプトを実装する
  • e.g. データベーススキーマからDTOを作成するスクリプトを実装する
  • 生成AIに任せてDTOを直接作成することも可能だが、求める制約が複雑な場合や生成対象が多い場合、作成ミスが発生しうる
• 実際のアプリケーション側のコードについては、用意したDTOを使ってもらうように生成AIに実装させる
  • 部品は用意しているので、こちらの意図した実装となりやすい

生成AIによるスクリプト生成例

プロンプト: postgresのCREATE TABLE文からdata classを生成するプラグインをbuildSrcにkotlinpoetで作成してください

import com.squareup.kotlinpoet.*
import net.sf.jsqlparser.parser.CCJSqlParserUtil
import net.sf.jsqlparser.statement.create.table.CreateTable
import org.gradle.api.DefaultTask
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.api.file.DirectoryProperty
import org.gradle.api.provider.Property
import org.gradle.api.tasks.*
import java.io.File

abstract class SqlToDataClassTask : DefaultTask() {
    @get:Input
    abstract val packageName: Property<String>

    @get:InputDirectory
    abstract val sqlInputDir: DirectoryProperty

    @get:OutputDirectory
    abstract val outputDir: DirectoryProperty

    private fun String.toClassName(): String {
        return split("_").joinToString("") { it.capitalize() }
    }

    private fun String.toPropertyName(): String {
        val parts = split("_")
        return parts[0] + parts.drop(1).joinToString("") { it.capitalize() }
    }

    private fun mapSqlTypeToKotlin(sqlType: String): TypeName {
        return when (sqlType.uppercase()) {
            "INTEGER", "INT", "SERIAL" -> Int::class.asTypeName()
            "BIGINT", "BIGSERIAL" -> Long::class.asTypeName()
            "TEXT", "VARCHAR", "CHAR", "CHARACTER VARYING" -> String::class.asTypeName()
            "BOOLEAN" -> Boolean::class.asTypeName()
            "NUMERIC", "DECIMAL" -> java.math.BigDecimal::class.asTypeName()
            "TIMESTAMP", "TIMESTAMP WITH TIME ZONE" -> java.time.OffsetDateTime::class.asTypeName()
            "DATE" -> java.time.LocalDate::class.asTypeName()
            "UUID" -> java.util.UUID::class.asTypeName()
            "JSONB", "JSON" -> String::class.asTypeName()
            else -> Any::class.asTypeName()
        }
    }

    @TaskAction
    fun execute() {
        val sqlDir = sqlInputDir.get().asFile
        val outputDirectory = outputDir.get().asFile
        outputDirectory.mkdirs()

        sqlDir.walkTopDown()
            .filter { it.extension == "sql" }
            .forEach { file ->
                val sql = file.readText()
                val statement = CCJSqlParserUtil.parse(sql)

                if (statement is CreateTable) {
                    val className = statement.table.name.toClassName()
                    val properties = statement.columnDefinitions.map { column ->
                        val propertyName = column.columnName.toPropertyName()
                        val type = mapSqlTypeToKotlin(column.colDataType.dataType)
                        val nullable = !column.columnSpecs?.any { it.toString().equals("NOT NULL", ignoreCase = true) } ?: true

                        PropertySpec.builder(propertyName, if (nullable) type.copy(nullable = true) else type)
                            .initializer(propertyName)
                            .build()
                    }

                    val kdoc = KDoc.builder()
                        .addParagraph("Generated from SQL table ${statement.table.name}")
                        .build()

                    val dataClass = TypeSpec.classBuilder(className)
                        .addKdoc(kdoc)
                        .addModifiers(KModifier.DATA)
                        .primaryConstructor(FunSpec.constructorBuilder()
                            .addParameters(properties.map {
                                ParameterSpec.builder(it.name, it.type).build()
                            })
                            .build())
                        .addProperties(properties)
                        .build()

                    val fileSpec = FileSpec.builder(packageName.get(), className)
                        .addType(dataClass)
                        .build()

                    fileSpec.writeTo(outputDirectory)
                }
            }
    }
}

class SqlToDataClassPlugin : Plugin<Project> {
    override fun apply(project: Project) {
        val extension = project.extensions.create("sqlToDataClass", SqlToDataClassExtension::class.java)

        project.tasks.register("generateDataClasses", SqlToDataClassTask::class.java) {
            it.packageName.set(extension.packageName)
            it.sqlInputDir.set(extension.sqlInputDir)
            it.outputDir.set(extension.outputDir)
        }
    }
}

abstract class SqlToDataClassExtension {
    abstract val packageName: Property<String>
    abstract val sqlInputDir: DirectoryProperty
    abstract val outputDir: DirectoryProperty
}

プロジェクトでのユースケース

1. RepositoryとMapperの生成
   1. エンティティによって追加でメソッドが必要となるので、初期構築用
2. OpenAPIのYAMLからのDTOとvalidationメソッド生成