improve companion object checking logic

google · Feb 28, 2024 · a726254 · a726254
1 parent c56396b
commit a726254
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/...-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt b/...-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt
@@ -33,7 +33,6 @@ import org.jetbrains.kotlin.analysis.api.components.buildTypeParameterType
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.descriptors.java.JavaVisibilities
 import org.jetbrains.kotlin.psi.KtClassOrObject
-import org.jetbrains.kotlin.psi.KtObjectDeclaration
 
 class KSClassDeclarationImpl private constructor(internal val ktClassOrObjectSymbol: KtClassOrObjectSymbol) :
     KSClassDeclaration,
@@ -99,7 +98,7 @@ class KSClassDeclarationImpl private constructor(internal val ktClassOrObjectSym
     }
 
     override val isCompanionObject: Boolean by lazy {
-        (ktClassOrObjectSymbol.psi as? KtObjectDeclaration)?.isCompanion() ?: false
+        ktClassOrObjectSymbol.classKind == KtClassKind.COMPANION_OBJECT
     }
 
     override fun getSealedSubclasses(): Sequence<KSClassDeclaration> {

diff --git a/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/ClassWithCompanionProcessor.kt b/test-utils/src/main/kotlin/com/google/devtools/ksp/processor/ClassWithCompanionProcessor.kt
@@ -17,6 +17,7 @@
 
 package com.google.devtools.ksp.processor
 
+import com.google.devtools.ksp.getClassDeclarationByName
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.symbol.KSAnnotated
 import com.google.devtools.ksp.symbol.KSClassDeclaration
@@ -29,6 +30,7 @@ class ClassWithCompanionProcessor : AbstractTestProcessor() {
 
     override fun process(resolver: Resolver): List<KSAnnotated> {
         resolver.getNewFiles().forEach { it.accept(CompanionVisitor(), Unit) }
+        resolver.getClassDeclarationByName("A")?.accept(CompanionVisitor(), Unit)
         return emptyList()
     }
 

diff --git a/test-utils/testData/api/companion.kt b/test-utils/testData/api/companion.kt
@@ -22,8 +22,17 @@
 // companion:false
 // obj:false
 // K:true
+// A:false
+// Companion:true
 // END
 
+// MODULE: lib
+// FILE: a.kt
+class A {
+    companion object {}
+}
+
+// MODULE: main(lib)
 class Foo {
     object companion {}
     object obj {}