Add EP to customize "Add Import" weights

Fix for https://youtrack.jetbrains.com/issue/KTIJ-31471
JetBrains · Oct 1, 2024 · b8a6875 · b8a6875
1 parent ab6b24b
commit b8a6875
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 6 deletions.
diff --git a/plugins/kotlin/code-insight/fixes-k2/resources/kotlin.code-insight.fixes.k2.xml b/plugins/kotlin/code-insight/fixes-k2/resources/kotlin.code-insight.fixes.k2.xml
@@ -2,6 +2,12 @@
   <dependencies>
     <module name="kotlin.refactorings.k2"/>
   </dependencies>
+  <extensionPoints>
+    <extensionPoint
+            qualifiedName="org.jetbrains.kotlin.expressionImportWeigherFactory"
+            interface="org.jetbrains.kotlin.idea.k2.codeinsight.fixes.imprt.ExpressionImportWeigherFactory"
+            dynamic="true"/>
+  </extensionPoints>
   <extensions defaultExtensionNs="org.jetbrains.kotlin">
     <codeinsight.quickfix.registrar implementation="org.jetbrains.kotlin.idea.k2.codeinsight.fixes.KotlinK2QuickFixRegistrar"/>
   </extensions>

diff --git a/...es-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/fixes/imprt/ExpressionImportWeigher.kt b/...es-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/fixes/imprt/ExpressionImportWeigher.kt
@@ -1,6 +1,7 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.k2.codeinsight.fixes.imprt
 
+import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.psi.PsiElement
 import com.intellij.util.applyIf
 import org.jetbrains.kotlin.analysis.api.KaSession
@@ -28,7 +29,15 @@ interface ExpressionImportWeigher {
 
     companion object {
         context(KaSession)
-        fun createWeigher(element: PsiElement?): ExpressionImportWeigher =
+        fun createWeigher(element: PsiElement?): ExpressionImportWeigher {
+            for (factory in ExpressionImportWeigherFactory.EP_NAME.extensionList) {
+                factory.createWeigher(element)?.let { return it }
+            }
+            return createDefaultWeigher(element)
+        }
+
+        context(KaSession)
+        private fun createDefaultWeigher(element: PsiElement?): ExpressionImportWeigher =
             when (element) {
                 is KtNameReferenceExpression -> CallExpressionImportWeigher(
                     token,
@@ -237,4 +246,14 @@ internal class OperatorExpressionImportWeigher(
 
         return weight
     }
-}
+}
+
+interface ExpressionImportWeigherFactory {
+    context(KaSession)
+    fun createWeigher(element: PsiElement?): ExpressionImportWeigher?
+
+    companion object {
+        internal val EP_NAME: ExtensionPointName<ExpressionImportWeigherFactory> =
+            ExtensionPointName.create("org.jetbrains.kotlin.expressionImportWeigherFactory")
+    }
+}
diff --git a/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/actions/ExpressionWeigher.kt b/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/actions/ExpressionWeigher.kt
@@ -1,6 +1,7 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.actions
 
+import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.analyze
 import org.jetbrains.kotlin.analysis.api.permissions.KaAllowAnalysisOnEdt
@@ -31,14 +32,21 @@ import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlin.types.typeUtil.replaceArgumentsWithStarProjections
 
 /**
- * Implementation in K2: [org.jetbrains.kotlin.idea.quickfix.importFix.ExpressionImportWeigher]
+ * Implementation in K2: [org.jetbrains.kotlin.idea.k2.codeinsight.fixes.imprt.ExpressionImportWeigher]
  */
-internal interface ExpressionWeigher {
+interface ExpressionWeigher {
 
     fun weigh(descriptor: DeclarationDescriptor): Int
 
     companion object {
-        fun createWeigher(element: PsiElement?): ExpressionWeigher =
+        internal fun createWeigher(element: PsiElement?): ExpressionWeigher {
+            for (factory in ExpressionWeigherFactory.EP_NAME.extensionList) {
+                factory.createWeigher(element)?.let { return it }
+            }
+            return createDefaultWeigher(element)
+        }
+
+        internal fun createDefaultWeigher(element: PsiElement?): ExpressionWeigher =
             when (element) {
                 is KtNameReferenceExpression -> CallExpressionWeigher(element)
                 is KtOperationReferenceExpression -> OperatorExpressionWeigher(element)
@@ -275,4 +283,13 @@ internal class OperatorExpressionWeigher(element: KtOperationReferenceExpression
         return weight
     }
 
-}
+}
+
+interface ExpressionWeigherFactory {
+    fun createWeigher(element: PsiElement?): ExpressionWeigher?
+
+    companion object {
+        internal val EP_NAME: ExtensionPointName<ExpressionWeigherFactory> =
+            ExtensionPointName.create("org.jetbrains.kotlin.expressionWeigherFactory")
+    }
+}
diff --git a/plugins/kotlin/plugin/k1/resources/META-INF/kotlin-core-fe10.xml b/plugins/kotlin/plugin/k1/resources/META-INF/kotlin-core-fe10.xml
@@ -18,6 +18,10 @@
             qualifiedName="org.jetbrains.kotlin.classImportFilter"
             interface="org.jetbrains.kotlin.idea.util.ClassImportFilter"
             dynamic="true"/>
+    <extensionPoint
+            qualifiedName="org.jetbrains.kotlin.expressionWeigherFactory"
+            interface="org.jetbrains.kotlin.idea.actions.ExpressionWeigherFactory"
+            dynamic="true"/>
   </extensionPoints>
 
   <extensions defaultExtensionNs="com.intellij">