自研IntelliJ插件如何修改Kotlin代码？PSI是否支持该操作？

完全可以用IntelliJ PSI操作Kotlin代码，JetBrains针对Kotlin提供了完整的PSI API，不用被Java相关的文档局限。下面结合你要实现的「给指定函数的when语句新增分支」需求，分享具体实现步骤：

核心步骤 #

1. 配置Kotlin PSI依赖 #

在你的插件项目中，确保引入Kotlin PSI相关依赖。如果是Gradle构建，在build.gradle（或build.gradle.kts）中添加：

dependencies {
    implementation "org.jetbrains.kotlin:kotlin-psi:$kotlinVersion"
    implementation "com.intellij:kotlin-plugin:$ideaVersion"
}

注意kotlinVersion和ideaVersion要与你开发插件对应的Android Studio版本匹配。

2. 定位目标When表达式 #

先找到目标Kotlin函数，再从中定位到要修改的when语句：

// 假设已获取到目标Kotlin文件的PsiFile对象
val ktFile = psiFile as KtFile

// 找到指定名称的函数（比如示例中的test函数）
val targetFunction = ktFile.declarations
    .filterIsInstance<KtNamedFunction>()
    .firstOrNull { it.name == "test" }

// 从函数体中筛选出when表达式
val whenExpression = targetFunction?.bodyExpression
    ?.statements
    ?.filterIsInstance<KtWhenExpression>()
    ?.firstOrNull()

3. 构造新的When分支 #

使用KtPsiFactory创建符合语法规范的新分支元素：

val psiFactory = KtPsiFactory(project)
// 生成新分支的PSI节点
val newBranch = psiFactory.createWhenEntry("3 -> NumberThree()")

4. 插入分支到When语句中 #

找到when语句的分支列表，将新分支插入到指定位置（比如示例中的最后一个分支之后）：

whenExpression?.let { whenExpr ->
    val entries = whenExpr.entries
    if (entries.isNotEmpty()) {
        // 所有PSI修改必须在WriteAction中执行，保证线程安全
        WriteAction.run<Throwable> {
            entries.last().parent.addAfter(newBranch, entries.last())
        }
    }
}

关键注意事项 #

• 所有PSI修改操作必须包裹在WriteAction中执行，这是IntelliJ平台的线程安全强制要求。
• KtPsiFactory能帮你生成任意合法的Kotlin代码元素，自动处理语法校验问题。
• 如果需要更精准定位（比如只修改针对变量num的when语句），可以通过whenExpression.subjectExpression判断目标变量是否匹配。

内容的提问来源于stack exchange，提问作者me.at.coding