KSP typealias fixes (#1956)

* Add support for KSTypeAlias.toClassName() * Set up resolveKSClassDeclaration that follows aliases * Use new resolving declaration check * Add test * Support aliases in KSType.toClassName() + fix nullability * Add toClassNameOrNull utility * Preserve aliases in toAnnotationSpec * Spotless * Changelog * More changelog + api * yay for simplicity
square · Aug 8, 2024 · fc64ac3 · fc64ac3
1 parent 193aea6
commit fc64ac3
Show file tree

Hide file tree

Showing 7 changed files with 100 additions and 40 deletions.
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -4,6 +4,10 @@ Change Log
 ## Unreleased
 
 * Fix: Enum classes that only have an init block now also generate the required semicolon (#1952)
+* Fix: Preserve typealiases in `KSAnnotation.toAnnotationSpec()`. (#1945)
+* Fix: Preserve nullability in `KSType.toClassName()`. (#1956)
+* New: Add `KSTypeAlias.toClassName()`. (#1956)
+* New: Add `KSType.toClassNameOrNull()`. (#1956)
 
 ## Version 1.18.1
 

diff --git a/interop/ksp/api/ksp.api b/interop/ksp/api/ksp.api
@@ -6,10 +6,12 @@ public final class com/squareup/kotlinpoet/ksp/AnnotationsKt {
 
 public final class com/squareup/kotlinpoet/ksp/KsClassDeclarationsKt {
 	public static final fun toClassName (Lcom/google/devtools/ksp/symbol/KSClassDeclaration;)Lcom/squareup/kotlinpoet/ClassName;
+	public static final fun toClassName (Lcom/google/devtools/ksp/symbol/KSTypeAlias;)Lcom/squareup/kotlinpoet/ClassName;
 }
 
 public final class com/squareup/kotlinpoet/ksp/KsTypesKt {
 	public static final fun toClassName (Lcom/google/devtools/ksp/symbol/KSType;)Lcom/squareup/kotlinpoet/ClassName;
+	public static final fun toClassNameOrNull (Lcom/google/devtools/ksp/symbol/KSType;)Lcom/squareup/kotlinpoet/ClassName;
 	public static final fun toTypeName (Lcom/google/devtools/ksp/symbol/KSType;Lcom/squareup/kotlinpoet/ksp/TypeParameterResolver;)Lcom/squareup/kotlinpoet/TypeName;
 	public static final fun toTypeName (Lcom/google/devtools/ksp/symbol/KSTypeArgument;Lcom/squareup/kotlinpoet/ksp/TypeParameterResolver;)Lcom/squareup/kotlinpoet/TypeName;
 	public static final fun toTypeName (Lcom/google/devtools/ksp/symbol/KSTypeReference;Lcom/squareup/kotlinpoet/ksp/TypeParameterResolver;)Lcom/squareup/kotlinpoet/TypeName;

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
@@ -21,31 +21,28 @@ import com.google.devtools.ksp.symbol.KSAnnotation
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSName
 import com.google.devtools.ksp.symbol.KSType
-import com.google.devtools.ksp.symbol.KSTypeAlias
 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.Companion.parameterizedBy
+import com.squareup.kotlinpoet.ParameterizedTypeName
 
 /**
  * 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 = 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 typeName = annotationType.resolve().toTypeName()
+
+  val builder = if (typeName is ClassName) {
+    AnnotationSpec.builder(typeName)
+  } else {
+    AnnotationSpec.builder(typeName as ParameterizedTypeName)
+  }
 
-  val params = (annotationType.resolve().declaration as KSClassDeclaration).primaryConstructor?.parameters.orEmpty()
+  val params = annotationType.resolve()
+    .resolveKSClassDeclaration()?.primaryConstructor?.parameters.orEmpty()
     .associateBy { it.name }
   useSiteTarget?.let { builder.useSiteTarget(it.kpAnalog) }
 
@@ -108,14 +105,6 @@ private val AnnotationUseSiteTarget.kpAnalog: UseSiteTarget
     AnnotationUseSiteTarget.DELEGATE -> UseSiteTarget.DELEGATE
   }
 
-internal fun KSType.unwrapTypeAlias(): KSType {
-  return if (this.declaration is KSTypeAlias) {
-    (this.declaration as KSTypeAlias).type.resolve()
-  } else {
-    this
-  }
-}
-
 private fun addValueToBlock(value: Any, member: CodeBlock.Builder, omitDefaultValues: Boolean) {
   when (value) {
     is List<*> -> {
@@ -140,14 +129,15 @@ private fun addValueToBlock(value: Any, member: CodeBlock.Builder, omitDefaultVa
     }
 
     is KSType -> {
-      val unwrapped = value.unwrapTypeAlias()
-      val isEnum = (unwrapped.declaration as KSClassDeclaration).classKind == ClassKind.ENUM_ENTRY
+      val declaration = value.resolveKSClassDeclaration() ?: error("Cannot resolve type of $value")
+      val isEnum = declaration.classKind == ClassKind.ENUM_ENTRY
       if (isEnum) {
-        val parent = unwrapped.declaration.parentDeclaration as KSClassDeclaration
-        val entry = unwrapped.declaration.simpleName.getShortName()
+        val parent = declaration.parentDeclaration?.resolveKSClassDeclaration()
+          ?: error("Could not resolve enclosing enum class of entry ${declaration.qualifiedName?.asString()}")
+        val entry = declaration.simpleName.getShortName()
         member.add("%T.%L", parent.toClassName(), entry)
       } else {
-        member.add("%T::class", unwrapped.toClassName())
+        member.add("%T::class", declaration.toClassName())
       }
     }
 

diff --git a/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/KsClassDeclarations.kt b/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/KsClassDeclarations.kt
@@ -16,9 +16,15 @@
 package com.squareup.kotlinpoet.ksp
 
 import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.google.devtools.ksp.symbol.KSTypeAlias
 import com.squareup.kotlinpoet.ClassName
 
 /** Returns the [ClassName] representation of this [KSClassDeclaration]. */
 public fun KSClassDeclaration.toClassName(): ClassName {
   return toClassNameInternal()
 }
+
+/** Returns the [ClassName] representation of this [KSTypeAlias]. */
+public fun KSTypeAlias.toClassName(): ClassName {
+  return toClassNameInternal()
+}
diff --git a/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/KsTypes.kt b/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/KsTypes.kt
@@ -36,20 +36,36 @@ import com.squareup.kotlinpoet.TypeVariableName
 import com.squareup.kotlinpoet.WildcardTypeName
 import com.squareup.kotlinpoet.tags.TypeAliasTag
 
-private fun KSType.requireNotErrorType() {
-  require(!isError) {
-    "Error type '$this' is not resolvable in the current round of processing."
-  }
-}
-
-/** Returns the [ClassName] representation of this [KSType] IFF it's a [KSClassDeclaration]. */
+/**
+ * Returns the [ClassName] representation of this [KSType] IFF it's a [KSClassDeclaration] or [KSTypeAlias].
+ */
 public fun KSType.toClassName(): ClassName {
   requireNotErrorType()
-  val decl = declaration
-  check(decl is KSClassDeclaration) {
-    "Declaration was not a KSClassDeclaration: $this"
+  check(arguments.isEmpty()) {
+    "KSType '$this' has type arguments, which are not supported for ClassName conversion. Use KSType.toTypeName()."
   }
-  return decl.toClassName()
+  return when (val decl = declaration) {
+    is KSClassDeclaration -> decl.toClassName()
+    is KSTypeAlias -> decl.toClassName()
+    is KSTypeParameter -> error("Cannot convert KSTypeParameter to ClassName: '$this'")
+    else -> error("Could not compute ClassName for '$this'")
+  }.copy(nullable = isMarkedNullable) as ClassName
+}
+
+/**
+ * Returns the [ClassName] representation of this [KSType] IFF it's a [KSClassDeclaration] or [KSTypeAlias].
+ *
+ * If it's unable to resolve to a [ClassName] for any reason, this returns null.
+ */
+public fun KSType.toClassNameOrNull(): ClassName? {
+  if (isError) return null
+  if (arguments.isNotEmpty()) return null
+  return when (val decl = declaration) {
+    is KSClassDeclaration -> decl.toClassName()
+    is KSTypeAlias -> decl.toClassName()
+    is KSTypeParameter -> null
+    else -> null
+  }?.let { it.copy(nullable = isMarkedNullable) as ClassName }
 }
 
 /**
@@ -73,7 +89,6 @@ internal fun KSType.toTypeName(
     is KSClassDeclaration -> {
       decl.toClassName().withTypeArguments(arguments.map { it.toTypeName(typeParamResolver) })
     }
-
     is KSTypeParameter -> typeParamResolver[decl.name.getShortName()]
     is KSTypeAlias -> {
       var typeAlias: KSTypeAlias = decl
@@ -107,11 +122,10 @@ internal fun KSType.toTypeName(
 
       val aliasArgs = typeArguments.map { it.toTypeName(typeParamResolver) }
 
-      decl.toClassNameInternal()
+      decl.toClassName()
         .withTypeArguments(aliasArgs)
         .copy(tags = mapOf(TypeAliasTag::class to TypeAliasTag(abbreviatedType)))
     }
-
     else -> error("Unsupported type: $declaration")
   }
 

diff --git a/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/utils.kt b/interop/ksp/src/main/kotlin/com/squareup/kotlinpoet/ksp/utils.kt
@@ -19,6 +19,9 @@ import com.google.devtools.ksp.isLocal
 import com.google.devtools.ksp.symbol.ClassKind
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSDeclaration
+import com.google.devtools.ksp.symbol.KSType
+import com.google.devtools.ksp.symbol.KSTypeAlias
+import com.google.devtools.ksp.symbol.KSTypeParameter
 import com.squareup.kotlinpoet.ClassName
 import com.squareup.kotlinpoet.LambdaTypeName
 import com.squareup.kotlinpoet.ParameterizedTypeName
@@ -76,3 +79,40 @@ internal fun KSDeclaration.toClassNameInternal(): ClassName {
     .split(".")
   return ClassName(pkgName, simpleNames)
 }
+
+internal fun KSType.requireNotErrorType() {
+  require(!isError) {
+    "Error type '$this' is not resolvable in the current round of processing."
+  }
+}
+
+/**
+ * Resolves the [KSClassDeclaration] for this type, including following typealiases as needed.
+ */
+internal fun KSType.resolveKSClassDeclaration(): KSClassDeclaration? {
+  requireNotErrorType()
+  return declaration.resolveKSClassDeclaration()
+}
+
+/**
+ * Resolves the [KSClassDeclaration] representation of this declaration, including following
+ * typealiases as needed.
+ *
+ * [KSTypeParameter] types will return null. If you expect one here, you should check the
+ * declaration directly.
+ */
+internal fun KSDeclaration.resolveKSClassDeclaration(): KSClassDeclaration? {
+  return when (val declaration = unwrapTypealiases()) {
+    is KSClassDeclaration -> declaration
+    is KSTypeParameter -> null
+    else -> error("Unexpected declaration type: $this")
+  }
+}
+
+/**
+ * Returns the resolved declaration following any typealiases.
+ */
+internal tailrec fun KSDeclaration.unwrapTypealiases(): KSDeclaration = when (this) {
+  is KSTypeAlias -> type.resolve().declaration.unwrapTypealiases()
+  else -> this
+}
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
@@ -925,7 +925,10 @@ class TestProcessorTest(private val useKsp2: Boolean) {
            typealias DaggerProvider<T> = @JvmSuppressWildcards Provider<T>
            interface SelectOptions
            interface SelectHandler<T>
+           annotation class SomeAnnotation
+           typealias AliasedAnnotation = SomeAnnotation
 
+           @AliasedAnnotation
            @ExampleAnnotation
            class Example(
              private val handlers: Map<Class<out SelectOptions>, DaggerProvider<SelectHandler<*>>>,
@@ -946,6 +949,7 @@ class TestProcessorTest(private val useKsp2: Boolean) {
         import java.lang.Class
         import kotlin.collections.Map
 
+        @AliasedAnnotation
         public class TestExample {
           private val handlers: Map<Class<out SelectOptions>, DaggerProvider<SelectHandler<*>>> = TODO()
         }