Remove blank lines after context receiver

facebook · Apr 30, 2023 · 1675a29 · 1675a29
1 parent 24b9234
commit 1675a29
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -95,6 +95,7 @@ import org.jetbrains.kotlin.psi.KtQualifiedExpression
 import org.jetbrains.kotlin.psi.KtReferenceExpression
 import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.KtScript
+import org.jetbrains.kotlin.psi.KtScriptInitializer
 import org.jetbrains.kotlin.psi.KtSecondaryConstructor
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
@@ -2432,6 +2433,7 @@ class KotlinInputAstVisitor(
   override fun visitScript(script: KtScript) {
     markForPartialFormat()
     var lastChildHadBlankLineBefore = false
+    var lastChildIsContextReceiver = false
     var first = true
     for (child in script.blockExpression.children) {
       if (child.text.isBlank()) {
@@ -2441,13 +2443,16 @@ class KotlinInputAstVisitor(
       val childGetsBlankLineBefore = child !is KtProperty
       if (first) {
         builder.blankLineWanted(OpsBuilder.BlankLineWanted.PRESERVE)
+      } else if (lastChildIsContextReceiver) {
+        builder.blankLineWanted(OpsBuilder.BlankLineWanted.NO)
       } else if (child !is PsiComment &&
           (childGetsBlankLineBefore || lastChildHadBlankLineBefore)) {
         builder.blankLineWanted(OpsBuilder.BlankLineWanted.YES)
       }
       visit(child)
       builder.guessToken(";")
       lastChildHadBlankLineBefore = childGetsBlankLineBefore
+      lastChildIsContextReceiver = child is KtScriptInitializer && child.firstChild?.firstChild?.firstChild?.text == "context"
       first = false
     }
     markForPartialFormat()

diff --git a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -6635,6 +6635,8 @@ class FormatterTest {
   fun `context receivers`() {
     val code =
         """
+      |context(Something)
+      |
       |class A {
       |  context(Logger, Raise<Error>)
       |
@@ -6651,6 +6653,7 @@ class FormatterTest {
 
     val expected =
         """
+      |context(Something)
       |class A {
       |  context(Logger, Raise<Error>)
       |  @SomeAnnotation

diff --git a/core/src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt b/core/src/test/java/com/facebook/ktfmt/format/TokenizerTest.kt
@@ -114,6 +114,7 @@ class TokenizerTest {
   @Test
   fun `Context receivers are parsed correctly`() {
     val code = """
+      |context(Something)
       |class A {
       |  context(Logger, Raise<Error>)
       |  fun test() {}
@@ -125,10 +126,10 @@ class TokenizerTest {
     file.accept(tokenizer)
 
     assertThat(tokenizer.toks.map { it.originalText })
-            .containsExactly("class", " ", "A", " ", "{", "\n", "  ", "context", "(", "Logger", ",", " ", "Raise", "<", "Error", ">", ")", "\n", "  ", "fun", " ", "test", "(", ")", " ", "{", "}", "\n", "}")
+            .containsExactly("context", "(", "Something", ")", "\n", "class", " ", "A", " ", "{", "\n", "  ", "context", "(", "Logger", ",", " ", "Raise", "<", "Error", ">", ")", "\n", "  ", "fun", " ", "test", "(", ")", " ", "{", "}", "\n", "}")
             .inOrder()
     assertThat(tokenizer.toks.map { it.index })
-            .containsExactly(0, -1, 1, -1, 2, -1, -1, 3, 4, 5, 6, -1, 7, 8, 9, 10, 11, -1, -1, 12, -1, 13, 14, 15, -1, 16, 17, -1, 18)
+            .containsExactly(0, 1, 2, 3, -1, 4, -1, 5, -1, 6, -1, -1, 7, 8, 9, 10, -1, 11, 12, 13, 14, 15, -1, -1, 16, -1, 17, 18, 19, -1, 20, 21, -1, 22)
             .inOrder()
   }
 }