Router: support all enumerators, including case

scalameta · Jun 18, 2024 · 503e238 · 503e238
1 parent b1172d4
commit 503e238
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/FormatOps.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
@@ -2879,4 +2879,9 @@ object FormatOps {
     case _ => None
   })
 
+  val EnumeratorAssignRhsOnLeft = new ExtractFromMeta(_.leftOwner match {
+    case x: Tree.WithBody with Enumerator => Some(x.body)
+    case _ => None
+  })
+
 }
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
@@ -2149,13 +2149,11 @@ class Router(formatOps: FormatOps) {
         Seq(Split(Space.orNL(!useNL), 0))
 
       // ForYield
-      case FormatToken(_: T.LeftArrow, _, _)
-          if leftOwner.is[Enumerator.Generator] =>
-        val enumerator = leftOwner.asInstanceOf[Enumerator.Generator]
-        getSplitsEnumerator(ft, enumerator.rhs)
-      case FormatToken(_: T.Equals, _, _) if leftOwner.is[Enumerator.Val] =>
-        val enumerator = leftOwner.asInstanceOf[Enumerator.Val]
-        getSplitsEnumerator(ft, enumerator.rhs)
+      case FormatToken(
+            _: T.LeftArrow | _: T.Equals,
+            _,
+            EnumeratorAssignRhsOnLeft(rhs),
+          ) => getSplitsEnumerator(ft, rhs)
 
       case FormatToken(
             kw @ (_: T.KwTry | _: T.KwCatch | _: T.KwFinally),