Fix ktlint-disable directive on package statement

README.md

@@ -511,22 +511,38 @@ See [Creating A Ruleset](#creating-a-ruleset).
 to produce the correct result (please report any such instances using [GitHub Issues](https://github.com/pinterest/ktlint/issues)).
 
 To disable a specific rule you'll need to turn on the verbose mode (`ktlint --verbose ...`). At the end of each line
-you'll see an error code. Use it as an argument for `ktlint-disable` directive (shown below).  
+you'll see the id of the rule. This id can be used as an argument for `ktlint-disable` directive (shown below).  
 
+A specific rule can be disabled for one specific line by adding the `ktlint-disable` directive, and the rule id as 
+follows:
 ```kotlin
-import package.* // ktlint-disable no-wildcard-imports
+import com.example.mypackage.* // ktlint-disable no-wildcard-imports
+```
 
+A rule can be disabled for a block of statements by adding the `ktlint-disable` directive, and the rule id as follows:
+```
 /* ktlint-disable no-wildcard-imports */
-import package.a.*
-import package.b.*
+import com.example.mypackage1.*
+import com.example.mypackage2.*
 /* ktlint-enable no-wildcard-imports */
 ```
+Note: the `ktlint-enable` directive above can be omitted in case the rule should be disabled for the remainder of the 
+file.
 
-To disable all checks:
+
+It is also possible to disable all checks for one single line, or a block of lines by not specifying the rule id after 
+the `ktlint-disable` directive:
 
 ```kotlin
-import package.* // ktlint-disable
+import com.example.mypackage.* // ktlint-disable
+
+/* ktlint-disable */
+import com.example.mypackage1.*
+import com.example.mypackage2.*
+/* ktlint-enable */
 ```
+Note: the `ktlint-enable` directive above can be omitted in case all rules should be disabled for the remainder of the 
+file.
 
 ### How do I globally disable a rule?
 See the [EditorConfig section](https://github.com/pinterest/ktlint#editorconfig) for details on how to use the `disabled_rules` property.

ktlint-core/src/main/kotlin/com/pinterest/ktlint/core/KtLint.kt

@@ -2,6 +2,7 @@ package com.pinterest.ktlint.core
 
 import com.pinterest.ktlint.core.api.EditorConfigProperties
 import com.pinterest.ktlint.core.api.FeatureInAlphaState
+import com.pinterest.ktlint.core.ast.isRoot
 import com.pinterest.ktlint.core.ast.visit
 import com.pinterest.ktlint.core.internal.EditorConfigGenerator
 import com.pinterest.ktlint.core.internal.EditorConfigLoader
@@ -88,16 +89,22 @@ public object KtLint {
         val errors = mutableListOf<LintError>()
 
         visitor(preparedCode.rootNode, params.ruleSets).invoke { node, rule, fqRuleId ->
+            if (node.isRoot() && rule is Rule.Modifier.InitializeRoot) {
+                // In case the rule has a root initializer then it should always be called regardless whether offset 0
+                // of the file is suppressed by a ktlint-disable directive.
+                rule.initializeRoot(node)
+            }
+
             // fixme: enforcing suppression based on node.startOffset is wrong
             // (not just because not all nodes are leaves but because rules are free to emit (and fix!) errors at any position)
             if (
-                !preparedCode.suppressedRegionLocator(node.startOffset, fqRuleId, node === preparedCode.rootNode)
+                !preparedCode.suppressedRegionLocator(node.startOffset, fqRuleId)
             ) {
                 try {
                     rule.visit(node, false) { offset, errorMessage, canBeAutoCorrected ->
                         // https://github.com/shyiko/ktlint/issues/158#issuecomment-462728189
                         if (node.startOffset != offset &&
-                            preparedCode.suppressedRegionLocator(offset, fqRuleId, node === preparedCode.rootNode)
+                            preparedCode.suppressedRegionLocator(offset, fqRuleId)
                         ) {
                             return@visit
                         }
@@ -297,10 +304,16 @@ public object KtLint {
         var mutated = false
         visitor(preparedCode.rootNode, params.ruleSets, concurrent = false)
             .invoke { node, rule, fqRuleId ->
+                if (node.isRoot() && rule is Rule.Modifier.InitializeRoot) {
+                    // In case the rule has a root initializer then it should always be called regardless whether offset
+                    // 0 of the file is suppressed by a ktlint-disable directive.
+                    rule.initializeRoot(node)
+                }
+
                 // fixme: enforcing suppression based on node.startOffset is wrong
                 // (not just because not all nodes are leaves but because rules are free to emit (and fix!) errors at any position)
                 if (
-                    !preparedCode.suppressedRegionLocator(node.startOffset, fqRuleId, node === preparedCode.rootNode)
+                    !preparedCode.suppressedRegionLocator(node.startOffset, fqRuleId)
                 ) {
                     try {
                         rule.visit(node, true) { _, _, canBeAutoCorrected ->
@@ -324,17 +337,23 @@ public object KtLint {
         if (tripped) {
             val errors = mutableListOf<Pair<LintError, Boolean>>()
             visitor(preparedCode.rootNode, params.ruleSets).invoke { node, rule, fqRuleId ->
+                if (node.isRoot() && rule is Rule.Modifier.InitializeRoot) {
+                    // In case the rule has a root initializer then it should always be called regardless whether offset 0
+                    // of the file is suppressed by a ktlint-disable directive.
+                    rule.initializeRoot(node)
+                }
+
                 // fixme: enforcing suppression based on node.startOffset is wrong
                 // (not just because not all nodes are leaves but because rules are free to emit (and fix!) errors at any position)
                 if (
-                    !preparedCode.suppressedRegionLocator(node.startOffset, fqRuleId, node === preparedCode.rootNode)
+                    !preparedCode.suppressedRegionLocator(node.startOffset, fqRuleId)
                 ) {
                     try {
                         rule.visit(node, false) { offset, errorMessage, canBeAutoCorrected ->
                             // https://github.com/shyiko/ktlint/issues/158#issuecomment-462728189
                             if (
                                 node.startOffset != offset &&
-                                preparedCode.suppressedRegionLocator(offset, fqRuleId, node === preparedCode.rootNode)
+                                preparedCode.suppressedRegionLocator(offset, fqRuleId)
                             ) {
                                 return@visit
                             }

ktlint-core/src/main/kotlin/com/pinterest/ktlint/core/Rule.kt

@@ -47,5 +47,13 @@ abstract class Rule(val id: String) {
          */
         interface RestrictToRootLast : RestrictToRoot
         interface Last
+        /**
+         * Any rule implementing this interface will be given the root ([FileASTNode]) node so that the rule is
+         * guaranteed to be initialized even in case the first line of the file contains a ktlint-disable directive
+         * for the rule.
+         */
+        interface InitializeRoot {
+            fun initializeRoot(node: ASTNode)
+        }
     }
 }

ktlint-core/src/main/kotlin/com/pinterest/ktlint/core/internal/SupressedRegionLocator.kt

@@ -1,24 +1,27 @@
 package com.pinterest.ktlint.core.internal
 
-import com.pinterest.ktlint.core.ast.prevLeaf
+import com.pinterest.ktlint.core.ast.ElementType
 import com.pinterest.ktlint.core.ast.visit
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.psi.PsiComment
+import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtAnnotationEntry
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.prevLeaf
+import org.jetbrains.kotlin.psi.psiUtil.prevLeafs
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 /**
  * Detects if given `ruleId` at given `offset` is suppressed.
  */
-internal typealias SuppressionLocator = (offset: Int, ruleId: String, isRoot: Boolean) -> Boolean
+internal typealias SuppressionLocator = (offset: Int, ruleId: String) -> Boolean
 
 /**
  * No suppression is detected. Always returns `false`.
  */
-internal val noSuppression: SuppressionLocator = { _, _, _ -> false }
+internal val noSuppression: SuppressionLocator = { _, _ -> false }
 
 private val suppressAnnotationRuleMap = mapOf(
     "RemoveCurlyBracesFromTemplate" to "string-template"
@@ -35,21 +38,13 @@ internal fun buildSuppressedRegionsLocator(
     val hintsList = collect(rootNode)
     return if (hintsList.isEmpty()) {
         noSuppression
-    } else { offset, ruleId, isRoot ->
-        if (isRoot) {
-            val h = hintsList[0]
-            h.range.last == 0 && (
-                h.disabledRules.isEmpty() ||
-                    h.disabledRules.contains(ruleId)
-                )
-        } else {
-            hintsList.any { hint ->
-                (
-                    hint.disabledRules.isEmpty() ||
-                        hint.disabledRules.contains(ruleId)
-                    ) &&
-                    hint.range.contains(offset)
-            }
+    } else { offset: Int, ruleId: String ->
+        hintsList.any { hint ->
+            (
+                hint.disabledRules.isEmpty() ||
+                    hint.disabledRules.contains(ruleId)
+                ) &&
+                hint.range.contains(offset)
         }
     }
 }
@@ -73,32 +68,14 @@ private fun collect(
             val text = node.getText()
             if (text.startsWith("//")) {
                 val commentText = text.removePrefix("//").trim()
-                parseHintArgs(commentText, "ktlint-disable")?.let { args ->
-                    val lineStart = (
-                        node.prevLeaf { it is PsiWhiteSpace && it.textContains('\n') } as
-                            PsiWhiteSpace?
-                        )?.let { it.node.startOffset + it.text.lastIndexOf('\n') + 1 } ?: 0
-                    result.add(SuppressionHint(IntRange(lineStart, node.startOffset), HashSet(args)))
-                }
+                node.createLineDisableSupressionHint(commentText)
+                    ?.let { suppressionHint -> result.add(suppressionHint) }
             } else {
                 val commentText = text.removePrefix("/*").removeSuffix("*/").trim()
-                parseHintArgs(commentText, "ktlint-disable")?.apply {
-                    open.add(SuppressionHint(IntRange(node.startOffset, node.startOffset), HashSet(this)))
-                }
-                    ?: parseHintArgs(commentText, "ktlint-enable")?.apply {
-                        // match open hint
-                        val disabledRules = HashSet(this)
-                        val openHintIndex = open.indexOfLast { it.disabledRules == disabledRules }
-                        if (openHintIndex != -1) {
-                            val openingHint = open.removeAt(openHintIndex)
-                            result.add(
-                                SuppressionHint(
-                                    IntRange(openingHint.range.first, node.startOffset),
-                                    disabledRules
-                                )
-                            )
-                        }
-                    }
+                node.createBlockDisableSuppressionHint(commentText)
+                    ?.let { suppressionHint -> open.add(suppressionHint) }
+                node.createBlockEnableSuppressionHint(commentText, open)
+                    ?.let { suppressionHint -> result.add(suppressionHint) }
             }
         }
         // Extract all Suppress annotations and create SuppressionHints
@@ -118,15 +95,50 @@ private fun collect(
     return result
 }
 
+private fun PsiElement.isDisabledOnPackageStatement(): Boolean =
+    prevLeafs.any { it.node.elementType == ElementType.PACKAGE_KEYWORD }
+
+private fun PsiElement.createLineDisableSupressionHint(commentText: String): SuppressionHint? {
+    return parseHintArgs(commentText, "ktlint-disable")
+        ?.let { hints ->
+            val rangeStartOffset = (
+                prevLeaf { it is PsiWhiteSpace && it.textContains('\n') } as
+                    PsiWhiteSpace?
+                )?.let { it.node.startOffset + it.text.lastIndexOf('\n') + 1 } ?: 0
+            SuppressionHint(IntRange(rangeStartOffset, startOffset), hints)
+        }
+}
+
+private fun PsiElement.createBlockDisableSuppressionHint(commentText: String): SuppressionHint? {
+    return parseHintArgs(commentText, "ktlint-disable")
+        ?.let { hints -> SuppressionHint(IntRange(node.startOffset, Int.MAX_VALUE), hints) }
+}
+
+private fun PsiElement.createBlockEnableSuppressionHint(commentText: String, open: java.util.ArrayList<SuppressionHint>): SuppressionHint? {
+    return parseHintArgs(commentText, "ktlint-enable")?.let { hints ->
+        // match open hint
+        val openHintIndex = open.indexOfLast { it.disabledRules == hints }
+        if (openHintIndex != -1) {
+            val openingHint = open.removeAt(openHintIndex)
+            SuppressionHint(
+                IntRange(openingHint.range.first, node.startOffset),
+                hints
+            )
+        } else {
+            null
+        }
+    }
+}
+
 private fun parseHintArgs(
     commentText: String,
     key: String
-): List<String>? {
+): HashSet<String>? {
     if (commentText.startsWith(key)) {
         val parsedComment = splitCommentBySpace(commentText)
         // assert exact match
         if (parsedComment[0] == key) {
-            return parsedComment.tail()
+            return HashSet(parsedComment.tail())
         }
     }
     return null

ktlint-core/src/test/kotlin/com/pinterest/ktlint/core/ErrorSuppressionTest.kt

@@ -10,29 +10,31 @@ import org.junit.Test
 
 class ErrorSuppressionTest {
 
-    @Test
-    fun testErrorSuppression() {
-        class NoWildcardImportsRule : Rule("no-wildcard-imports") {
-            override fun visit(
-                node: ASTNode,
-                autoCorrect: Boolean,
-                emit: (offset: Int, errorMessage: String, corrected: Boolean) -> Unit
-            ) {
-                if (node is LeafPsiElement && node.textMatches("*") && node.isPartOf(ElementType.IMPORT_DIRECTIVE)) {
-                    emit(node.startOffset, "Wildcard import", false)
-                }
+    class NoWildcardImportsRule : Rule("no-wildcard-imports") {
+        override fun visit(
+            node: ASTNode,
+            autoCorrect: Boolean,
+            emit: (offset: Int, errorMessage: String, corrected: Boolean) -> Unit
+        ) {
+            if (node is LeafPsiElement && node.textMatches("*") && node.isPartOf(ElementType.IMPORT_DIRECTIVE)) {
+                emit(node.startOffset, "Wildcard import", false)
             }
         }
-        fun lint(text: String) =
-            ArrayList<LintError>().apply {
-                KtLint.lint(
-                    KtLint.Params(
-                        text = text,
-                        ruleSets = listOf(RuleSet("standard", NoWildcardImportsRule())),
-                        cb = { e, _ -> add(e) }
-                    )
+    }
+
+    fun lint(text: String) =
+        ArrayList<LintError>().apply {
+            KtLint.lint(
+                KtLint.Params(
+                    text = text,
+                    ruleSets = listOf(RuleSet("standard", NoWildcardImportsRule())),
+                    cb = { e, _ -> add(e) }
                 )
-            }
+            )
+        }
+
+    @Test
+    fun testErrorSuppressionDisableAllOnImport() {
         assertThat(
             lint(
                 """
@@ -45,6 +47,10 @@ class ErrorSuppressionTest {
                 LintError(2, 10, "no-wildcard-imports", "Wildcard import")
             )
         )
+    }
+
+    @Test
+    fun testErrorSuppressionDisableRuleOnImport() {
         assertThat(
             lint(
                 """
@@ -57,6 +63,10 @@ class ErrorSuppressionTest {
                 LintError(2, 10, "no-wildcard-imports", "Wildcard import")
             )
         )
+    }
+
+    @Test
+    fun testErrorSuppressionDisableAllInBlock() {
         assertThat(
             lint(
                 """
@@ -72,6 +82,10 @@ class ErrorSuppressionTest {
                 LintError(5, 10, "no-wildcard-imports", "Wildcard import")
             )
         )
+    }
+
+    @Test
+    fun testErrorSuppressionDisableRuleInBlock() {
         assertThat(
             lint(
                 """