Support shebang for KTS files

Summary: This was crashing in JavaInput, so let's just remove it and add it back at the end.

Reviewed By: cortinico

Differential Revision: D36987246

fbshipit-source-id: e404d213d9ba5a916caaa63cba850066d1d77436

2 files changed, 21 insertions, 3 deletions

diff --git a/core/src/main/java/com/facebook/ktfmt/format/Formatter.kt b/core/src/main/java/com/facebook/ktfmt/format/Formatter.kt
index 3bf2e2a..1fa90f7 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/Formatter.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/Formatter.kt
@@ -78,9 +78,15 @@ object Formatter {
   @JvmStatic
   @Throws(FormatterException::class, ParseError::class)
   fun format(options: FormattingOptions, code: String): String {
-    checkEscapeSequences(code)
+    val (shebang, kotlinCode) =
+        if (code.startsWith("#!")) {
+          code.split("\n".toRegex(), limit = 2)
+        } else {
+          listOf("", code)
+        }
+    checkEscapeSequences(kotlinCode)
 
-    val lfCode = StringUtilRt.convertLineSeparators(code)
+    val lfCode = StringUtilRt.convertLineSeparators(kotlinCode)
     val sortedImports = sortedAndDistinctImports(lfCode)
     val pretty = prettyPrint(sortedImports, options, "\n")
     val noRedundantElements =
@@ -89,7 +95,9 @@ object Formatter {
         } catch (e: ParseError) {
           throw IllegalStateException("Failed to re-parse code after pretty printing:\n $pretty", e)
         }
-    return prettyPrint(noRedundantElements, options, Newlines.guessLineSeparator(code)!!)
+    val prettyCode =
+        prettyPrint(noRedundantElements, options, Newlines.guessLineSeparator(kotlinCode)!!)
+    return if (shebang.isNotEmpty()) shebang + "\n" + prettyCode else prettyCode
   }
 
   /** prettyPrint reflows 'code' using google-java-format's engine. */
diff --git a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
index b403a01..bf57e02 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -50,6 +50,16 @@ class FormatterTest {
         |""".trimMargin())
 
   @Test
+  fun `support script (kts) files with a shebang`() =
+      assertFormatted(
+          """
+        |#!/usr/bin/env kscript
+        |package foo
+        |
+        |println("Called")
+        |""".trimMargin())
+
+  @Test
   fun `call chains`() =
       assertFormatted(
           """