Add support for annotation type arguments in KSP (#1889)

square · Apr 18, 2024 · 55ec0ce · 55ec0ce
1 parent d73fac9
commit 55ec0ce
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 19 deletions.
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -6,6 +6,7 @@ Change Log
 * New: Supertype list wraps to one-per-line if the primary constructor spans multiple lines (#1866).
 * New: Extract `MemberSpecHolder` interface for constructs that can hold `PropertySpec`s and `FunSpec`s and their builders (#1877).
 * New: `joinToCode` variant which operates on any type, but requires a transform lambda to convert each element into a `CodeBlock`.
+* New: Support annotation type arguments in `KSAnnotation.toAnnotationSpec()` (#1889).
 * Fix: Prevent name clashes between a function in class and a function call in current scope (#1850).
 * Fix: Fix extension function imports (#1814).
 * Fix: Omit implicit modifiers on FileSpec.scriptBuilder (#1813).

diff --git a/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/Annotations.kt b/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/Annotations.kt
@@ -26,31 +26,39 @@ import com.squareup.kotlinpoet.AnnotationSpec
 import com.squareup.kotlinpoet.AnnotationSpec.UseSiteTarget
 import com.squareup.kotlinpoet.ClassName
 import com.squareup.kotlinpoet.CodeBlock
-import com.squareup.kotlinpoet.ParameterizedTypeName
+import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
 
 /**
  * Returns an [AnnotationSpec] representation of this [KSAnnotation] instance.
  * @param omitDefaultValues omit defining default values when `true`
  */
 @JvmOverloads
 public fun KSAnnotation.toAnnotationSpec(omitDefaultValues: Boolean = false): AnnotationSpec {
-  val builder = when (val type = annotationType.resolve().unwrapTypeAlias().toTypeName()) {
-    is ClassName -> AnnotationSpec.builder(type)
-    is ParameterizedTypeName -> AnnotationSpec.builder(type)
-    else -> error("This is never possible.")
-  }
+  val builder = annotationType.resolve().unwrapTypeAlias().toClassName()
+    .let { className ->
+      val typeArgs = annotationType.element?.typeArguments.orEmpty()
+        .map { it.toTypeName() }
+      if (typeArgs.isEmpty()) {
+        AnnotationSpec.builder(className)
+      } else {
+        AnnotationSpec.builder(className.parameterizedBy(typeArgs))
+      }
+    }
+
   val params = (annotationType.resolve().declaration as KSClassDeclaration).primaryConstructor?.parameters.orEmpty()
     .associateBy { it.name }
   useSiteTarget?.let { builder.useSiteTarget(it.kpAnalog) }
-  // TODO support type params once they're exposed https://github.com/google/ksp/issues/753
+
   var varargValues: List<*>? = null
   for (argument in arguments) {
     val value = argument.value ?: continue
     val name = argument.name!!.getShortName()
     val type = params[argument.name]
     if (omitDefaultValues) {
       val defaultValue = this.defaultArguments.firstOrNull { it.name?.asString() == name }?.value
-      if (isDefaultValue(value, defaultValue)) { continue }
+      if (isDefaultValue(value, defaultValue)) {
+        continue
+      }
     }
     if (type?.isVararg == true) {
       // Wait to add varargs to end.
@@ -87,17 +95,18 @@ private fun isDefaultValue(value: Any?, defaultValue: Any?): Boolean {
   return value == defaultValue
 }
 
-private val AnnotationUseSiteTarget.kpAnalog: UseSiteTarget get() = when (this) {
-  AnnotationUseSiteTarget.FILE -> UseSiteTarget.FILE
-  AnnotationUseSiteTarget.PROPERTY -> UseSiteTarget.PROPERTY
-  AnnotationUseSiteTarget.FIELD -> UseSiteTarget.FIELD
-  AnnotationUseSiteTarget.GET -> UseSiteTarget.GET
-  AnnotationUseSiteTarget.SET -> UseSiteTarget.SET
-  AnnotationUseSiteTarget.RECEIVER -> UseSiteTarget.RECEIVER
-  AnnotationUseSiteTarget.PARAM -> UseSiteTarget.PARAM
-  AnnotationUseSiteTarget.SETPARAM -> UseSiteTarget.SETPARAM
-  AnnotationUseSiteTarget.DELEGATE -> UseSiteTarget.DELEGATE
-}
+private val AnnotationUseSiteTarget.kpAnalog: UseSiteTarget
+  get() = when (this) {
+    AnnotationUseSiteTarget.FILE -> UseSiteTarget.FILE
+    AnnotationUseSiteTarget.PROPERTY -> UseSiteTarget.PROPERTY
+    AnnotationUseSiteTarget.FIELD -> UseSiteTarget.FIELD
+    AnnotationUseSiteTarget.GET -> UseSiteTarget.GET
+    AnnotationUseSiteTarget.SET -> UseSiteTarget.SET
+    AnnotationUseSiteTarget.RECEIVER -> UseSiteTarget.RECEIVER
+    AnnotationUseSiteTarget.PARAM -> UseSiteTarget.PARAM
+    AnnotationUseSiteTarget.SETPARAM -> UseSiteTarget.SETPARAM
+    AnnotationUseSiteTarget.DELEGATE -> UseSiteTarget.DELEGATE
+  }
 
 internal fun KSType.unwrapTypeAlias(): KSType {
   return if (this.declaration is KSTypeAlias) {
@@ -129,6 +138,7 @@ private fun addValueToBlock(value: Any, member: CodeBlock.Builder, omitDefaultVa
       }
       member.add("⇤⇤)")
     }
+
     is KSType -> {
       val unwrapped = value.unwrapTypeAlias()
       val isEnum = (unwrapped.declaration as KSClassDeclaration).classKind == ClassKind.ENUM_ENTRY
@@ -140,12 +150,14 @@ private fun addValueToBlock(value: Any, member: CodeBlock.Builder, omitDefaultVa
         member.add("%T::class", unwrapped.toClassName())
       }
     }
+
     is KSName ->
       member.add(
         "%T.%L",
         ClassName.bestGuess(value.getQualifier()),
         value.getShortName(),
       )
+
     is KSAnnotation -> member.add("%L", value.toAnnotationSpec(omitDefaultValues))
     else -> member.add(memberForValue(value))
   }

diff --git a/...rocessor/src/main/kotlin/com/squareup/kotlinpoet/ksp/test/processor/exampleAnnotations.kt b/...rocessor/src/main/kotlin/com/squareup/kotlinpoet/ksp/test/processor/exampleAnnotations.kt
@@ -81,3 +81,4 @@ enum class AnnotationEnumValue {
 }
 
 annotation class AnnotationWithVararg(val simpleArg: Int, vararg val args: String)
+annotation class AnnotationWithTypeArgs<T, R>
diff --git a/...processor/src/test/kotlin/com/squareup/kotlinpoet/ksp/test/processor/TestProcessorTest.kt b/...processor/src/test/kotlin/com/squareup/kotlinpoet/ksp/test/processor/TestProcessorTest.kt
@@ -819,6 +819,45 @@ class TestProcessorTest {
     )
   }
 
+  @Test
+  fun typeArgs() {
+    val compilation = prepareCompilation(
+      kotlin(
+        "Example.kt",
+        """
+           package test
+
+           import com.squareup.kotlinpoet.ksp.test.processor.ExampleAnnotation
+           import com.squareup.kotlinpoet.ksp.test.processor.AnnotationWithTypeArgs
+
+           @ExampleAnnotation
+           @AnnotationWithTypeArgs<String, List<Int>>
+           class Example
+           """,
+      ),
+    )
+
+    val result = compilation.compile()
+    assertThat(result.exitCode).isEqualTo(KotlinCompilation.ExitCode.OK)
+    val generatedFileText = File(compilation.kspSourcesDir, "kotlin/test/TestExample.kt")
+      .readText()
+
+    assertThat(generatedFileText).isEqualTo(
+      """
+        package test
+
+        import com.squareup.kotlinpoet.ksp.test.processor.AnnotationWithTypeArgs
+        import kotlin.Int
+        import kotlin.String
+        import kotlin.collections.List
+
+        @AnnotationWithTypeArgs<String, List<Int>>
+        public class TestExample
+
+      """.trimIndent(),
+    )
+  }
+
   private fun prepareCompilation(vararg sourceFiles: SourceFile): KotlinCompilation {
     return KotlinCompilation()
       .apply {