Router: add indent to for-yield with comments

scalameta · Dec 9, 2021 · d0dc8b6 · d0dc8b6
1 parent e9971ab
commit d0dc8b6
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala
@@ -2133,15 +2133,18 @@ class Router(formatOps: FormatOps) {
 
       // Term.ForYield
       case FormatToken(T.KwYield(), _, _) if leftOwner.is[Term.ForYield] =>
+        val lastToken =
+          getLastToken(leftOwner.asInstanceOf[Term.ForYield].body)
+        val indent = Indent(style.indent.main, lastToken, ExpiresOn.After)
         if (style.newlines.avoidAfterYield && !rightOwner.is[Term.If]) {
-          Seq(Split(Space, 0))
+          val nextFt = nextNonCommentSameLine(formatToken)
+          val noIndent = nextFt.eq(formatToken) || nextFt.noBreak
+          Seq(Split(Space, 0).withIndent(indent, noIndent))
         } else {
-          val lastToken =
-            getLastToken(leftOwner.asInstanceOf[Term.ForYield].body)
           Seq(
             // Either everything fits in one line or break on =>
             Split(Space, 0).withSingleLineNoOptimal(lastToken),
-            Split(Newline, 1).withIndent(style.indent.main, lastToken, After)
+            Split(Newline, 1).withIndent(indent)
           )
         }
 

diff --git a/scalafmt-tests/src/test/resources/test/IndentYieldKeyword.stat b/scalafmt-tests/src/test/resources/test/IndentYieldKeyword.stat
@@ -125,8 +125,8 @@ object a {
     a <- b
     c <- d
   } yield // c
-  a
+    a
   for { a <- b } yield // c
-  a +
-    b
+    a +
+      b
 }