Kotlin 元编程之 KotlinPoet

Posted 2023-01-19 川峰

在 KSP 中默认生成代码的方式是通过CodeGenerator 创建文件流后以字符串拼接的方式来生成代码，对于简单的demo还好，但是对于实际生产项目中要生成的代码可能会十分复杂，如果还是自己手动去拼接，可能非常的繁琐，累死人不说，还非常容易出错，比如说少拼接了一个标点符号，可能需要排查半天。实际生产项目中使用的最多的就是由 JakeWharton 大神所写的著名的开源库 JavaPoet（很有诗意的名字，翻译过来叫Java诗人）使用该库可通过方便的函数进行拼接，减少出错。

KotlinPoet 是对应 JavaPoet 的 Kotlin 版本，同样是由square开发的，它可以用来很方便的生成 Kotlin 代码。

本文介绍 KotlinPoet 的使用，包括但不限于其官网文档中的内容，你也可以直接参考其官方文档：https://square.github.io/kotlinpoet/

配置

在ksp模块的build.gradle中添加KotlinPoet的依赖：

dependencies 
    implementation 'com.squareup:kotlinpoet:1.12.0' 

对应版本可以在Github上的KotlinPoet官网上查找。

简单使用

例如：

val greeterClass = ClassName("com.example.generated", "Greeter")
val fileSpec = FileSpec.builder("com.example.generated", "HelloWorld")
    .addType(
        TypeSpec.classBuilder(greeterClass)
            .primaryConstructor(
                FunSpec.constructorBuilder()
                    .addParameter("name", String::class)
                    .build()
            )
            .addProperty(
                PropertySpec.builder("name", String::class)
                    .initializer("name")
                    .build()
            )
            .addFunction(
                FunSpec.builder("greet")
                    .addStatement("println(%P)", "Hello, \\$name")
                    .build()
            )
            .build()
    )
    .addFunction(
        FunSpec.builder("main")
            .addParameter("args", String::class, KModifier.VARARG)
            .addStatement("%T(args[0]).greet()", greeterClass)
            .build()
    )
    .build()

fileSpec.writeTo(System.out)

这会生成一个包含如下代码的HelloWorld.kt文件：

package com.example.generated

import kotlin.String
import kotlin.Unit

public class Greeter(
  public val name: String,
) 
  public fun greet(): Unit 
    println("""Hello, $name""")
  


public fun main(vararg args: String): Unit 
  Greeter(args[0]).greet()

是不是很简单，跟直接拼接的方式相比，可读性很好，而且更加安全。

KotlinPoet 根据不同的使用用途提供了不同的开箱即用的类：

生成目标	使用对象
Kotlin 文件	FileSpec，可以调用其addType、addFunction、addImport、addCode、addProperty等来生成文件内容
类、接口和对象	TypeSpec，可以调用其addModifiers、addFunctions、addProperty等来生成主体内容
函数和构造函数	FunSpec，可以调用其 addModifiers、addParameters、addStatement、addCode等来生成函数内容
参数	ParameterSpec
属性	PropertySpec
注解	AnnotationSpec
类型别名	TypeAliasSpec

addCode

但是方法和构造函数的主体在 KotlinPoet 中没有建模，没有表达式类、语句类或语法树节点。KotlinPoet 可通过调用 addCode 方法传入一个字符串模板作为代码块生成方式，可以利用 Kotlin 的多行字符串使它看起来更漂亮：

val main = FunSpec.builder("main")
  .addCode("""
    |var total = 0
    |for (i in 0 until 10) 
    |    total += i
    |
    |""".trimMargin())
  .build()

这样会生成如下代码：

fun main() 
  var total = 0
  for (i in 0 until 10) 
    total += i
  

ControlFlow

通过 addStatement 配合 beginControlFlow和endControlFlow可以进行更加灵活的流程控制代码生成：

private fun computeRange(name: String, from: Int, to: Int, op: String): FunSpec 
  return FunSpec.builder(name)
    .returns(Int::class)
    .addStatement("var result = 1")
    .beginControlFlow("for (i in $from until $to)")
    .addStatement("result = result $op i")
    .endControlFlow()
    .addStatement("return result")
    .build()

例如当调用 computeRange("computeRange", 1, 100, "*")时，会生成以下代码：

public fun computeRange(): Int 
  var result = 1
  for (i in 1 until 100) 
    result = result * i
  
  return result

%S 代表字符串

当使用字符串模板的方式生成代码时，使用%S代表一个String，它会完成包装引号和转义，例如：

fun main(args: Array<String>) 
  val helloWorld = TypeSpec.classBuilder("HelloWorld")
    .addFunction(whatsMyNameYo("slimShady"))
    .addFunction(whatsMyNameYo("eminem"))
    .addFunction(whatsMyNameYo("marshallMathers"))
    .build()

  val kotlinFile = FileSpec.builder("com.example.helloworld", "HelloWorld")
    .addType(helloWorld)
    .build()

  kotlinFile.writeTo(System.out)


private fun whatsMyNameYo(name: String): FunSpec 
  return FunSpec.builder(name)
    .returns(String::class)
    .addStatement("return %S", name)
    .build()

这会生成以下代码：

class HelloWorld 
  fun slimShady(): String = "slimShady"

  fun eminem(): String = "eminem"

  fun marshallMathers(): String = "marshallMathers"

使用%S会自动加上双引号。

%P 用于字符串模板

%S还会自动处理美元符号 ( $) 的转义，以避免无意中创建字符串模板导致无法在生成的代码中编译：

val stringWithADollar = "Your total is " + "$" + "50"
val funSpec = FunSpec.builder("printTotal")
  .returns(String::class)
  .addStatement("return %S", stringWithADollar)
  .build()

这会产生：

fun printTotal(): String = "Your total is $'$'50"

如果调用printTotal()函数，就会输出Your total is $50，可以看到美元符号被自动转义了，这很好，但是如果需要在拼接字符串模板时， $用于引用变量，而不转义美元符号，请使用 %P ：

private fun stringTemplate(): FunSpec
    val stringWithADollar = "Your total is " + "\\$amount"
    return FunSpec.builder("printTotal")
        .addParameter("amount", String::class)
        .returns(String::class)
        .addStatement("return %P", stringWithADollar)
        .build()

这会产生：

public fun printTotal(amount: String): String = """Your total is $amount"""

这样就是动态输出amount变量的值了。

CodeBlock

还可以将CodeBlocks 用作 %P 的参数，这在需要在字符串模板中引用可导入类型或成员时非常方便：

val file = FileSpec.builder("com.example", "Digits")
  .addFunction(
    FunSpec.builder("print")
      .addParameter("digits", IntArray::class)
      .addStatement("println(%P)", buildCodeBlock 
        val contentToString = MemberName("kotlin.collections", "contentToString")
        add("These are the digits: \\$digits.%M()", contentToString)
      )
      .build()
  )
  .build()
println(file)

上面的代码片段将产生以下输出，会正确的处理导入：

package com.example

import kotlin.IntArray
import kotlin.collections.contentToString

fun print(digits: IntArray) 
  println("""These are the digits: $digits.contentToString()""")

%T 引用类型自动导入

KotlinPoet 对类型有丰富的内置支持，包括 import 语句的自动生成。仅用于%T引用类型：

val today = FunSpec.builder("today")
  .returns(Date::class)
  .addStatement("return %T()", Date::class)
  .build()

val helloWorld = TypeSpec.classBuilder("HelloWorld")
  .addFunction(today)
  .build()

val kotlinFile = FileSpec.builder("com.example.helloworld", "HelloWorld")
  .addType(helloWorld)
  .build()

kotlinFile.writeTo(System.out)

这会生成以下.kt文件，其中包含必要内容的import：

package com.example.helloworld

import java.util.Date

class HelloWorld 
  fun today(): Date = Date()

ClassName 用于构建Class类型

上面我们通过Date::class引用了一个我们在编写生成代码时恰好可用的类。但是我们也可以引用一个在编写生成代码时还不存在的类：

val hoverboard = ClassName("com.mattel", "Hoverboard")

val tomorrow = FunSpec.builder("tomorrow")
  .returns(hoverboard)
  .addStatement("return %T()", hoverboard)
  .build()

这会生成以下代码，那个还不存在的类也被导入了：

package com.example.helloworld

import com.mattel.Hoverboard

class HelloWorld 
  fun tomorrow(): Hoverboard = Hoverboard()

由于类型非常重要，在使用 KotlinPoet 时会经常需要到 ClassName。它可以识别任何声明的类。声明类型只是 Kotlin 丰富类型系统的开始：我们还有数组、参数化类型、通配符类型、lambda 类型和类型变量。KotlinPoet 具有用于构建以下各项的类：

import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy

val hoverboard = ClassName("com.mattel", "Hoverboard")
val list = ClassName("kotlin.collections", "List")
val arrayList = ClassName("kotlin.collections", "ArrayList")
val listOfHoverboards = list.parameterizedBy(hoverboard)
val arrayListOfHoverboards = arrayList.parameterizedBy(hoverboard)

val thing = ClassName("com.misc", "Thing")
val array = ClassName("kotlin", "Array")
val producerArrayOfThings = array.parameterizedBy(WildcardTypeName.producerOf(thing))

val beyond = FunSpec.builder("beyond")
  .returns(listOfHoverboards)
  .addStatement("val result = %T()", arrayListOfHoverboards)
  .addStatement("result += %T()", hoverboard)
  .addStatement("result += %T()", hoverboard)
  .addStatement("result += %T()", hoverboard)
  .addStatement("return result")
  .build()

val printThings = FunSpec.builder("printThings")
  .addParameter("things", producerArrayOfThings)
  .addStatement("println(things)")
  .build()

这会生成以下代码，KotlinPoet 将分解每种类型并在可能的情况下将其导入：

package com.example.helloworld

import com.mattel.Hoverboard
import com.misc.Thing
import kotlin.Array
import kotlin.collections.ArrayList
import kotlin.collections.List

class HelloWorld 
  fun beyond(): List<Hoverboard> 
    val result = ArrayList<Hoverboard>()
    result += Hoverboard()
    result += Hoverboard()
    result += Hoverboard()
    return result
  

  fun printThings(things: Array<out Thing>) 
    println(things)
  

可空类型

KotlinPoet 支持可空类型。要将一个 TypeName 转换成可为 null 的对应项，请使用copy(nullable = true)方法:

val name = PropertySpec.builder("name", String::class.asTypeName().copy(nullable = true))
    .mutable()
    .addModifiers(KModifier.PRIVATE)
    .initializer("null")
    .build()
    
val address = PropertySpec.builder("address", String::class)
    .addModifiers(KModifier.PRIVATE)
    .initializer("%S", "china")
    .build()

TypeSpec.classBuilder("HelloWorld")
    .addProperty(name)
    .addProperty(address)
    //.addProperty("address", String::class, KModifier.PRIVATE)
    .build()

这会生成以下代码:

class HelloWorld 
  private var name: String? = null

  private val address: String = "china"

%M 引用 MemberName 成员

与ClassName类似，KotlinPoet 有一个特殊的成员占位符（函数和属性），当代码需要访问顶级成员和在对象内部声明的成员时，它会派上用场。用%M引用成员时，需要传递一个MemberName实例作为占位符的参数，KotlinPoet 将自动处理导入：

package com.squareup.tacos

class Taco 
fun createTaco() = Taco() 
val Taco.isVegan: Boolean
    get() = false

val createTaco = MemberName("com.squareup.tacos", "createTaco")
val isVegan = MemberName("com.squareup.tacos", "isVegan")
val file = FileSpec.builder("com.squareup.example", "TacoTest")
  .addFunction(
    FunSpec.builder("main")
      .addStatement("val taco = %M()", createTaco)
      .addStatement("println(taco.%M)", isVegan)
      .build()
  )
  .build()
println(file)

这会生成以下文件：

package com.squareup.example

import com.squareup.tacos.createTaco
import com.squareup.tacos.isVegan

fun main() 
  val taco = createTaco()
  println(taco.isVegan)

如您所见，%M