remove psi from visitKtlintSuppressionInAnnotation

pinterest · Dec 10, 2024 · a05595f · a05595f
1 parent 2562354
commit a05595f
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 8 deletions.
diff --git a/ktlint-rule-engine-core/api/ktlint-rule-engine-core.api b/ktlint-rule-engine-core/api/ktlint-rule-engine-core.api
@@ -56,7 +56,8 @@ public final class com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionKt
 	public static synthetic fun prevLeaf$default (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;ZILjava/lang/Object;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
 	public static final fun prevSibling (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lkotlin/jvm/functions/Function1;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
 	public static synthetic fun prevSibling$default (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
-	public static final fun recursiveChildren (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Lkotlin/sequences/Sequence;
+	public static final fun recursiveChildren (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Z)Lkotlin/sequences/Sequence;
+	public static synthetic fun recursiveChildren$default (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;ZILjava/lang/Object;)Lkotlin/sequences/Sequence;
 	public static final fun remove (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)V
 	public static final fun replaceWith (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)V
 	public static final fun upsertWhitespaceAfterMe (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Ljava/lang/String;)V

diff --git a/...engine-core/src/main/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtension.kt b/...engine-core/src/main/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtension.kt
@@ -230,8 +230,11 @@ public fun ASTNode.isPartOfComment(): Boolean =
 
 public fun ASTNode.children(): Sequence<ASTNode> = generateSequence(firstChildNode) { node -> node.treeNext }
 
-public fun ASTNode.recursiveChildren(): Sequence<ASTNode> =
+public fun ASTNode.recursiveChildren(includeSelf: Boolean = false): Sequence<ASTNode> =
     sequence {
+        if (includeSelf) {
+            yield(this@recursiveChildren)
+        }
         children().forEach {
             yield(it)
             yieldAll(it.recursiveChildren())

diff --git a/.../src/main/kotlin/com/pinterest/ktlint/rule/engine/internal/rules/KtlintSuppressionRule.kt b/.../src/main/kotlin/com/pinterest/ktlint/rule/engine/internal/rules/KtlintSuppressionRule.kt
@@ -17,6 +17,7 @@ import com.pinterest.ktlint.rule.engine.core.api.nextLeaf
 import com.pinterest.ktlint.rule.engine.core.api.nextSibling
 import com.pinterest.ktlint.rule.engine.core.api.parent
 import com.pinterest.ktlint.rule.engine.core.api.prevLeaf
+import com.pinterest.ktlint.rule.engine.core.api.recursiveChildren
 import com.pinterest.ktlint.rule.engine.core.api.remove
 import com.pinterest.ktlint.rule.engine.core.api.replaceWith
 import com.pinterest.ktlint.rule.engine.internal.KTLINT_SUPPRESSION_ID_ALL_RULES
@@ -42,9 +43,9 @@ import org.jetbrains.kotlin.psi.KtScriptInitializer
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
 import org.jetbrains.kotlin.psi.KtValueArgument
 import org.jetbrains.kotlin.psi.KtValueArgumentList
-import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
 import org.jetbrains.kotlin.psi.psiUtil.siblings
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.utils.addToStdlib.applyIf
 
 /**
@@ -105,10 +106,10 @@ public class KtlintSuppressionRule(
         emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> AutocorrectDecision,
     ) {
         node
-            .psi
-            .findDescendantOfType<KtLiteralStringTemplateEntry>()
-            ?.node
-            ?.let { literalStringTemplateEntry ->
+            .recursiveChildren(includeSelf = true)
+            .firstOrNull {
+                it.elementType == ElementType.LITERAL_STRING_TEMPLATE_ENTRY
+            }?.let { literalStringTemplateEntry ->
                 val prefixedSuppression =
                     literalStringTemplateEntry
                         .text
@@ -123,7 +124,9 @@ public class KtlintSuppressionRule(
                         .ifAutocorrectAllowed {
                             node
                                 .createLiteralStringTemplateEntry(prefixedSuppression)
-                                ?.let { literalStringTemplateEntry.replaceWith(it) }
+                                ?.let {
+                                    literalStringTemplateEntry.replaceWith(it)
+                                }
                         }
                 }
             }