Merge pull request #1030 from massimosiani/fix-stripmargin

make stripMargin similar to strings
typelevel · Jan 20, 2024 · f947177 · f947177
2 parents e74de91 + b853e4d
commit f947177
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 9 deletions.
diff --git a/modules/core/shared/src/main/scala/Fragment.scala b/modules/core/shared/src/main/scala/Fragment.scala
@@ -58,11 +58,18 @@ final case class Fragment[A](
   def stripMargin: Fragment[A] = stripMargin('|')
 
   def stripMargin(marginChar: Char): Fragment[A] = {
-    val ps = parts.map {
+    val head = parts.headOption
+    val tail = parts.tail
+    val ps = head.map {
       _.bimap(
-        _.stripMargin(marginChar).replaceAll("\n", " "),
-        _.map(_.stripMargin(marginChar).replaceAll("\n", " "))
+        _.stripMargin(marginChar),
+        _.map(_.stripMargin(marginChar))
       )
+    }.toList ++ tail.map {
+      _.bimap(
+        str => str.takeWhile(_ != '\n') + str.dropWhile(_ != '\n').stripMargin(marginChar),
+        _.map(str => str.takeWhile(_ != '\n') + str.dropWhile(_ != '\n').stripMargin(marginChar))
+  )
     }
     Fragment(ps, encoder, origin)
   }

diff --git a/modules/tests/shared/src/test/scala/FragmentTest.scala b/modules/tests/shared/src/test/scala/FragmentTest.scala
@@ -83,15 +83,26 @@ class FragmentTest extends SkunkTest {
 
   pureTest("stripMargin") {
     val f = sql"""select
-    |$int4
-    |""".stripMargin
-    f.sql.trim == sql"select $int4".sql
+              |$int4
+              |""".stripMargin
+    f.sql == """select
+               |$1
+               |""".stripMargin
   }
 
   pureTest("stripMargin with char") {
     val f = sql"""select
-    ^$int4
-    ^""".stripMargin('^')
-    f.sql.trim == sql"select $int4".sql
+                 ^$int4
+                 ^""".stripMargin('^')
+    f.sql == """select
+               ^$1
+               ^""".stripMargin('^')
+  }
+
+  pureTest("stripMargin respects intermediate pipes") {
+    val f = sql"""|select ${text} || 'foo'
+                  |2""".stripMargin
+    f.sql == """select $1 || 'foo'
+               |2""".stripMargin
   }
 }