如何修改Kotlin的PSI？IntelliJ插件开发技术求助

需求明确 #

需要开发IntelliJ IDEA插件，让IDE解析Kotlin接口时，自动将以下代码：

interface TestInterface {
    var testProperty: Int
}

识别为包含实现自身的伴生对象的版本：

interface TestInterface {
    var testProperty: Int

    companion object: TestInterface {
        override var testProperty: Int
            get() {
                TODO("Do something")
            }
            set(value) {
                TODO("Do something")
            }
    }
}

为什么kapt无法实现 #

kapt是Kotlin编译期注解处理器，核心作用是生成物理源文件，而你的需求是IDE层面的动态PSI扩展——让IDE在代码分析时“虚拟”识别出伴生对象，不修改实际文件，这不属于kapt的处理范畴，因此方向错误。

正确实现方案：基于IntelliJ Kotlin PSI扩展 #

使用KtPsiAugmentProvider实现动态PSI元素添加，这是IntelliJ Kotlin插件用于扩展代码洞察的标准方式。

1. 实现PsiAugmentProvider #

import com.intellij.openapi.project.Project
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import org.jetbrains.kotlin.psi.*
import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject

class InterfaceCompanionAugmentProvider : KtPsiAugmentProvider() {
    override fun <Psi : PsiElement> collectAugments(
        element: PsiElement,
        type: Class<Psi>,
        context: PsiFile
    ): List<Psi> {
        // 仅处理Kotlin接口
        val ktInterface = element as? KtClass ?: return emptyList()
        if (!ktInterface.isInterface()) return emptyList()
        // 避免重复生成伴生对象
        if (ktInterface.companionObjects.isNotEmpty()) return emptyList()

        val companionObject = generateCompanionObject(ktInterface.project, ktInterface) ?: return emptyList()
        @Suppress("UNCHECKED_CAST")
        return if (type.isInstance(companionObject)) listOf(companionObject as Psi) else emptyList()
    }

    private fun generateCompanionObject(project: Project, ktInterface: KtClass): KtObjectDeclaration? {
        val interfaceName = ktInterface.name ?: return null
        val factory = KtPsiFactory(project)

        // 构建伴生对象基础代码
        val companionBaseText = """
            companion object : $interfaceName {
                ${ktInterface.properties.joinToString("\n") { prop ->
                    "override var ${prop.name}: ${prop.typeReference?.text ?: "Unit"}"
                }}
            }
        """.trimIndent()

        val objectDeclaration = factory.createObjectDeclaration(companionBaseText)
        // 为每个属性添加自定义getter/setter
        objectDeclaration.properties.forEach { prop ->
            val getter = factory.createPropertyGetter("TODO(\"Do something\")")
            val setter = factory.createPropertySetter("value", "TODO(\"Do something\")")
            prop.add(getter)
            prop.add(setter)
        }
        return objectDeclaration
    }
}

2. 注册扩展 #

在插件的plugin.xml中添加以下配置，让IntelliJ识别这个扩展：

<extensions defaultExtensionNs="org.jetbrains.kotlin">
    <psiAugmentProvider implementation="io.github.sgpublic.exspplugin.core.kt.InterfaceCompanionAugmentProvider"/>
</extensions>

关键注意点 #

• 必须使用KtPsiAugmentProvider而非直接修改PSI：直接修改PSI会破坏IDE的代码一致性，而AugmentProvider是官方提供的动态扩展入口，确保代码分析、补全、跳转等功能正常。
• 校验重复生成：检查接口是否已有伴生对象，避免重复添加。
• 使用KtPsiFactory构建PSI元素：手动拼接文本容易出现语法错误，通过工厂类生成的PSI元素能保证符合Kotlin语法规范。

内容的提问来源于stack exchange，提问作者Madray Haven