Allow adding suffix to LineCommentCreator (#285)

sbt · Apr 21, 2022 · 01c2c35 · 01c2c35
1 parent 5fc984b
commit 01c2c35
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/src/main/scala/de/heikoseeberger/sbtheader/CommentStyle.scala b/src/main/scala/de/heikoseeberger/sbtheader/CommentStyle.scala
@@ -94,16 +94,22 @@ object TwirlStyleFramedBlockCommentCreator extends CommentCreator {
   private def stars(count: Int) = "*" * count
 }
 
-final class LineCommentCreator(linePrefix: String) extends CommentCreator {
+final class LineCommentCreator(linePrefix: String, lineSuffix: String = "") extends CommentCreator {
 
   override def apply(text: String, existingText: Option[String]): String = {
     def prependWithLinePrefix(s: String) =
       s match {
         case ""   => if (linePrefix.trim.nonEmpty) linePrefix else ""
         case line => s"$linePrefix $line"
       }
+    def appendLineSuffix(s: String) =
+      s match {
+        case "" => ""
+        case line => // Only add the suffix when a prefix was added before
+          if (linePrefix.trim.nonEmpty && lineSuffix.trim.nonEmpty) s"$line $lineSuffix" else line
+      }
 
-    text.linesIterator.map(prependWithLinePrefix).mkString(NewLine)
+    text.linesIterator.map(prependWithLinePrefix).map(appendLineSuffix).mkString(NewLine)
   }
 }