Refactor CLI argument parsing (#467)

Summary:
Working towards https://github.com/facebook/ktfmt/issues/391.

I wanted to open this to get feedback on the approach.

I could have chosen to use exceptions, but I find sealed classes much more composable and I already anticipate adding another class to the `ParseResult` hierarchy to represent a situation where a message needs to be shown to the user but isn't an error (`--help` is the obvious example).

Closes https://github.com/facebook/ktfmt/issues/465 more or less as a side-effect.

Pull Request resolved: https://github.com/facebook/ktfmt/pull/467

Reviewed By: strulovich

Differential Revision: D58136233

Pulled By: hick209

fbshipit-source-id: 85662e7b5f19195c4bb4c9c28caa02ae224ede52

3 files changed, 134 insertions, 103 deletions

diff --git a/core/src/main/java/com/facebook/ktfmt/cli/Main.kt b/core/src/main/java/com/facebook/ktfmt/cli/Main.kt
index 5f46eb9..fba648b 100644
--- a/core/src/main/java/com/facebook/ktfmt/cli/Main.kt
+++ b/core/src/main/java/com/facebook/ktfmt/cli/Main.kt
@@ -36,7 +36,7 @@ class Main(
     private val input: InputStream,
     private val out: PrintStream,
     private val err: PrintStream,
-    args: Array<String>
+    private val inputArgs: Array<String>
 ) {
   companion object {
     @JvmStatic
@@ -69,9 +69,13 @@ class Main(
     }
   }
 
-  private val parsedArgs: ParsedArgs = ParsedArgs.processArgs(err, args)
-
   fun run(): Int {
+    val processArgs = ParsedArgs.processArgs(inputArgs)
+    val parsedArgs =
+        when (processArgs) {
+          is ParseResult.Ok -> processArgs.parsedValue
+          is ParseResult.Error -> exitFatal(processArgs.errorMessage, 1)
+        }
     if (parsedArgs.fileNames.isEmpty()) {
       err.println(
           "Usage: ktfmt [--dropbox-style | --google-style | --kotlinlang-style] [--dry-run] [--set-exit-if-changed] [--stdin-name=<name>] [--do-not-remove-unused-imports] File1.kt File2.kt ...")
@@ -81,7 +85,7 @@ class Main(
 
     if (parsedArgs.fileNames.size == 1 && parsedArgs.fileNames[0] == "-") {
       return try {
-        val alreadyFormatted = format(null)
+        val alreadyFormatted = format(null, parsedArgs)
         if (!alreadyFormatted && parsedArgs.setExitIfChanged) 1 else 0
       } catch (e: Exception) {
         1
@@ -107,7 +111,7 @@ class Main(
     val retval = AtomicInteger(0)
     files.parallelStream().forEach {
       try {
-        if (!format(it) && parsedArgs.setExitIfChanged) {
+        if (!format(it, parsedArgs) && parsedArgs.setExitIfChanged) {
           retval.set(1)
         }
       } catch (e: Exception) {
@@ -126,17 +130,17 @@ class Main(
    * @param file The file to format. If null, the code is read from <stdin>.
    * @return true iff input is valid and already formatted.
    */
-  private fun format(file: File?): Boolean {
-    val fileName = file?.toString() ?: parsedArgs.stdinName ?: "<stdin>"
+  private fun format(file: File?, args: ParsedArgs): Boolean {
+    val fileName = file?.toString() ?: args.stdinName ?: "<stdin>"
     try {
       val bytes = if (file == null) input else FileInputStream(file)
       val code = BufferedReader(InputStreamReader(bytes, UTF_8)).readText()
-      val formattedCode = Formatter.format(parsedArgs.formattingOptions, code)
+      val formattedCode = Formatter.format(args.formattingOptions, code)
       val alreadyFormatted = code == formattedCode
 
       // stdin
       if (file == null) {
-        if (parsedArgs.dryRun) {
+        if (args.dryRun) {
           if (!alreadyFormatted) {
             out.println(fileName)
           }
@@ -146,7 +150,7 @@ class Main(
         return alreadyFormatted
       }
 
-      if (parsedArgs.dryRun) {
+      if (args.dryRun) {
         if (!alreadyFormatted) {
           out.println(fileName)
         }
@@ -173,4 +177,15 @@ class Main(
       throw e
     }
   }
+
+  /**
+   * Finishes the process with result `returnCode`.
+   *
+   * **WARNING**: If you call this method, this is the last that will happen and no code after it
+   * will be executed.
+   */
+  private fun exitFatal(message: String, returnCode: Int): Nothing {
+    err.println(message)
+    exitProcess(returnCode)
+  }
 }
diff --git a/core/src/main/java/com/facebook/ktfmt/cli/ParsedArgs.kt b/core/src/main/java/com/facebook/ktfmt/cli/ParsedArgs.kt
index e0009f7..0a5f7bd 100644
--- a/core/src/main/java/com/facebook/ktfmt/cli/ParsedArgs.kt
+++ b/core/src/main/java/com/facebook/ktfmt/cli/ParsedArgs.kt
@@ -19,7 +19,6 @@ package com.facebook.ktfmt.cli
 import com.facebook.ktfmt.format.Formatter
 import com.facebook.ktfmt.format.FormattingOptions
 import java.io.File
-import java.io.PrintStream
 import java.nio.charset.StandardCharsets.UTF_8
 
 /** ParsedArgs holds the arguments passed to ktfmt on the command-line, after parsing. */
@@ -39,16 +38,16 @@ data class ParsedArgs(
 ) {
   companion object {
 
-    fun processArgs(err: PrintStream, args: Array<String>): ParsedArgs {
+    fun processArgs(args: Array<String>): ParseResult {
       if (args.size == 1 && args[0].startsWith("@")) {
-        return parseOptions(err, File(args[0].substring(1)).readLines(UTF_8).toTypedArray())
+        return parseOptions(File(args[0].substring(1)).readLines(UTF_8).toTypedArray())
       } else {
-        return parseOptions(err, args)
+        return parseOptions(args)
       }
     }
 
     /** parseOptions parses command-line arguments passed to ktfmt. */
-    fun parseOptions(err: PrintStream, args: Array<String>): ParsedArgs {
+    fun parseOptions(args: Array<out String>): ParseResult {
       val fileNames = mutableListOf<String>()
       var formattingOptions = FormattingOptions()
       var dryRun = false
@@ -64,29 +63,36 @@ data class ParsedArgs(
           arg == "--dry-run" || arg == "-n" -> dryRun = true
           arg == "--set-exit-if-changed" -> setExitIfChanged = true
           arg == "--do-not-remove-unused-imports" -> removeUnusedImports = false
-          arg.startsWith("--stdin-name") -> stdinName = parseKeyValueArg(err, "--stdin-name", arg)
-          arg.startsWith("--") -> err.println("Unexpected option: $arg")
-          arg.startsWith("@") -> err.println("Unexpected option: $arg")
+          arg.startsWith("--stdin-name=") ->
+              stdinName =
+                  parseKeyValueArg("--stdin-name", arg)
+                      ?: return ParseResult.Error(
+                          "Found option '${arg}', expected '${"--stdin-name"}=<value>'")
+          arg.startsWith("--") -> return ParseResult.Error("Unexpected option: $arg")
+          arg.startsWith("@") -> return ParseResult.Error("Unexpected option: $arg")
           else -> fileNames.add(arg)
         }
       }
 
-      return ParsedArgs(
-          fileNames,
-          formattingOptions.copy(removeUnusedImports = removeUnusedImports),
-          dryRun,
-          setExitIfChanged,
-          stdinName,
-      )
+      return ParseResult.Ok(
+          ParsedArgs(
+              fileNames,
+              formattingOptions.copy(removeUnusedImports = removeUnusedImports),
+              dryRun,
+              setExitIfChanged,
+              stdinName,
+          ))
     }
 
-    private fun parseKeyValueArg(err: PrintStream, key: String, arg: String): String? {
+    private fun parseKeyValueArg(key: String, arg: String): String? {
       val parts = arg.split('=', limit = 2)
-      if (parts[0] != key || parts.size != 2) {
-        err.println("Found option '${arg}', expected '${key}=<value>'")
-        return null
-      }
-      return parts[1]
+      return parts[1].takeIf { parts[0] == key || parts.size == 2 }
     }
   }
 }
+
+sealed interface ParseResult {
+  data class Ok(val parsedValue: ParsedArgs) : ParseResult
+
+  data class Error(val errorMessage: String) : ParseResult
+}
diff --git a/core/src/test/java/com/facebook/ktfmt/cli/ParsedArgsTest.kt b/core/src/test/java/com/facebook/ktfmt/cli/ParsedArgsTest.kt
index 9f0919b..467bb41 100644
--- a/core/src/test/java/com/facebook/ktfmt/cli/ParsedArgsTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/cli/ParsedArgsTest.kt
@@ -19,9 +19,7 @@ package com.facebook.ktfmt.cli
 import com.facebook.ktfmt.format.Formatter
 import com.facebook.ktfmt.format.FormattingOptions
 import com.google.common.truth.Truth.assertThat
-import java.io.ByteArrayOutputStream
 import java.io.FileNotFoundException
-import java.io.PrintStream
 import kotlin.io.path.createTempDirectory
 import kotlin.test.assertFailsWith
 import org.junit.After
@@ -41,143 +39,106 @@ class ParsedArgsTest {
   }
 
   @Test
-  fun `files to format are returned and unknown flags are reported`() {
-    val (parsed, out) = parseTestOptions("foo.kt", "--unknown")
-
-    assertThat(parsed.fileNames).containsExactly("foo.kt")
-    assertThat(out.toString()).isEqualTo("Unexpected option: --unknown\n")
+  fun `unknown flags return an error`() {
+    val result = ParsedArgs.parseOptions(arrayOf("--unknown"))
+    assertThat(result).isInstanceOf(ParseResult.Error::class.java)
   }
 
   @Test
-  fun `files to format are returned and flags starting with @ are reported`() {
-    val (parsed, out) = parseTestOptions("foo.kt", "@unknown")
-
-    assertThat(parsed.fileNames).containsExactly("foo.kt")
-    assertThat(out.toString()).isEqualTo("Unexpected option: @unknown\n")
+  fun `unknown flags starting with '@' return an error`() {
+    val result = ParsedArgs.parseOptions(arrayOf("@unknown"))
+    assertThat(result).isInstanceOf(ParseResult.Error::class.java)
   }
 
   @Test
   fun `parseOptions uses default values when args are empty`() {
-    val (parsed, _) = parseTestOptions("foo.kt")
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("foo.kt")))
 
     val formattingOptions = parsed.formattingOptions
-    assertThat(formattingOptions.style).isEqualTo(FormattingOptions.Style.FACEBOOK)
-    assertThat(formattingOptions.maxWidth).isEqualTo(100)
-    assertThat(formattingOptions.blockIndent).isEqualTo(2)
-    assertThat(formattingOptions.continuationIndent).isEqualTo(4)
-    assertThat(formattingOptions.removeUnusedImports).isTrue()
-    assertThat(formattingOptions.debuggingPrintOpsAfterFormatting).isFalse()
 
-    assertThat(parsed.dryRun).isFalse()
-    assertThat(parsed.setExitIfChanged).isFalse()
-    assertThat(parsed.stdinName).isNull()
+    val defaultFormattingOptions = FormattingOptions()
+    assertThat(formattingOptions).isEqualTo(defaultFormattingOptions)
   }
 
   @Test
-  fun `parseOptions recognizes --dropbox-style and rejects unknown flags`() {
-    val (parsed, out) = parseTestOptions("--dropbox-style", "foo.kt", "--unknown")
-
-    assertThat(parsed.fileNames).containsExactly("foo.kt")
-    assertThat(parsed.formattingOptions.blockIndent).isEqualTo(4)
-    assertThat(parsed.formattingOptions.continuationIndent).isEqualTo(4)
-    assertThat(out.toString()).isEqualTo("Unexpected option: --unknown\n")
+  fun `parseOptions recognizes --dropbox-style`() {
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("--dropbox-style", "foo.kt")))
+    assertThat(parsed.formattingOptions).isEqualTo(Formatter.DROPBOX_FORMAT)
   }
 
   @Test
   fun `parseOptions recognizes --google-style`() {
-    val (parsed, _) = parseTestOptions("--google-style", "foo.kt")
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("--google-style", "foo.kt")))
     assertThat(parsed.formattingOptions).isEqualTo(Formatter.GOOGLE_FORMAT)
   }
 
   @Test
   fun `parseOptions recognizes --dry-run`() {
-    val (parsed, _) = parseTestOptions("--dry-run", "foo.kt")
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("--dry-run", "foo.kt")))
     assertThat(parsed.dryRun).isTrue()
   }
 
   @Test
   fun `parseOptions recognizes -n as --dry-run`() {
-    val (parsed, _) = parseTestOptions("-n", "foo.kt")
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("-n", "foo.kt")))
     assertThat(parsed.dryRun).isTrue()
   }
 
   @Test
   fun `parseOptions recognizes --set-exit-if-changed`() {
-    val (parsed, _) = parseTestOptions("--set-exit-if-changed", "foo.kt")
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("--set-exit-if-changed", "foo.kt")))
     assertThat(parsed.setExitIfChanged).isTrue()
   }
 
   @Test
   fun `parseOptions defaults to removing imports`() {
-    val (parsed, _) = parseTestOptions("foo.kt")
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("foo.kt")))
     assertThat(parsed.formattingOptions.removeUnusedImports).isTrue()
   }
 
   @Test
   fun `parseOptions recognizes --do-not-remove-unused-imports to removing imports`() {
-    val (parsed, _) = parseTestOptions("--do-not-remove-unused-imports", "foo.kt")
-    assertThat(parsed.formattingOptions.removeUnusedImports).isFalse()
-  }
-
-  @Test
-  fun `parseOptions handles dropbox style and --do-not-remove-unused-imports`() {
-    val (parsed, _) =
-        parseTestOptions("--do-not-remove-unused-imports", "--dropbox-style", "foo.kt")
+    val parsed =
+        assertSucceeds(ParsedArgs.parseOptions(arrayOf("--do-not-remove-unused-imports", "foo.kt")))
     assertThat(parsed.formattingOptions.removeUnusedImports).isFalse()
-    assertThat(parsed.formattingOptions.style).isEqualTo(FormattingOptions.Style.DROPBOX)
   }
 
   @Test
-  fun `parseOptions handles google style and --do-not-remove-unused-imports`() {
-    val (parsed, _) = parseTestOptions("--do-not-remove-unused-imports", "--google-style", "foo.kt")
-    assertThat(parsed.formattingOptions.removeUnusedImports).isFalse()
-    assertThat(parsed.formattingOptions.style).isEqualTo(FormattingOptions.Style.GOOGLE)
-  }
-
-  @Test
-  fun `parseOptions --stdin-name`() {
-    val (parsed, _) = parseTestOptions("--stdin-name=my/foo.kt")
+  fun `parseOptions recognizes --stdin-name`() {
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("--stdin-name=my/foo.kt")))
     assertThat(parsed.stdinName).isEqualTo("my/foo.kt")
   }
 
   @Test
-  fun `parseOptions --stdin-name with empty value`() {
-    val (parsed, _) = parseTestOptions("--stdin-name=")
+  fun `parseOptions accepts --stdin-name with empty value`() {
+    val parsed = assertSucceeds(ParsedArgs.parseOptions(arrayOf("--stdin-name=")))
     assertThat(parsed.stdinName).isEqualTo("")
   }
 
   @Test
   fun `parseOptions --stdin-name without value`() {
-    val (parsed, out) = parseTestOptions("--stdin-name")
-    assertThat(out).isEqualTo("Found option '--stdin-name', expected '--stdin-name=<value>'\n")
-    assertThat(parsed.stdinName).isNull()
-  }
-
-  @Test
-  fun `parseOptions --stdin-name prefix`() {
-    val (parsed, out) = parseTestOptions("--stdin-namea")
-    assertThat(out).isEqualTo("Found option '--stdin-namea', expected '--stdin-name=<value>'\n")
-    assertThat(parsed.stdinName).isNull()
+    val parseResult = ParsedArgs.parseOptions(arrayOf("--stdin-name"))
+    assertThat(parseResult).isInstanceOf(ParseResult.Error::class.java)
   }
 
   @Test
   fun `processArgs use the @file option with non existing file`() {
-    val out = ByteArrayOutputStream()
-
     val e =
         assertFailsWith<FileNotFoundException> {
-          ParsedArgs.processArgs(PrintStream(out), arrayOf("@non-existing-file"))
+          ParsedArgs.processArgs(arrayOf("@non-existing-file"))
         }
     assertThat(e.message).contains("non-existing-file (No such file or directory)")
   }
 
   @Test
   fun `processArgs use the @file option with file containing arguments`() {
-    val out = ByteArrayOutputStream()
     val file = root.resolve("existing-file")
     file.writeText("--google-style\n--dry-run\n--set-exit-if-changed\nFile1.kt\nFile2.kt\n")
 
-    val parsed = ParsedArgs.processArgs(PrintStream(out), arrayOf("@" + file.absolutePath))
+    val result = ParsedArgs.processArgs(arrayOf("@" + file.absolutePath))
+    assertThat(result).isInstanceOf(ParseResult.Ok::class.java)
+
+    val parsed = (result as ParseResult.Ok).parsedValue
 
     assertThat(parsed.formattingOptions).isEqualTo(Formatter.GOOGLE_FORMAT)
     assertThat(parsed.dryRun).isTrue()
@@ -185,8 +146,57 @@ class ParsedArgsTest {
     assertThat(parsed.fileNames).containsExactlyElementsIn(listOf("File1.kt", "File2.kt"))
   }
 
-  private fun parseTestOptions(vararg args: String): Pair<ParsedArgs, String> {
-    val out = ByteArrayOutputStream()
-    return Pair(ParsedArgs.parseOptions(PrintStream(out), arrayOf(*args)), out.toString())
+  @Test
+  fun `parses multiple args successfully`() {
+    val testResult =
+        ParsedArgs.parseOptions(
+            arrayOf("--google-style", "--dry-run", "--set-exit-if-changed", "File.kt"),
+        )
+    assertThat(testResult)
+        .isEqualTo(
+            parseResultOk(
+                fileNames = listOf("File.kt"),
+                formattingOptions = Formatter.GOOGLE_FORMAT,
+                dryRun = true,
+                setExitIfChanged = true,
+            ))
+  }
+
+  @Test
+  fun `last style in args wins`() {
+    val testResult =
+        ParsedArgs.parseOptions(arrayOf<String>("--google-style", "--dropbox-style", "File.kt"))
+    assertThat(testResult)
+        .isEqualTo(
+            parseResultOk(
+                fileNames = listOf("File.kt"),
+                formattingOptions = Formatter.DROPBOX_FORMAT,
+            ))
+  }
+
+  @Test
+  fun `error when parsing multiple args and one is unknown`() {
+    val testResult =
+        ParsedArgs.parseOptions(arrayOf<String>("@unknown", "--google-style", "File.kt"))
+    assertThat(testResult).isEqualTo(ParseResult.Error("Unexpected option: @unknown"))
+  }
+
+  private fun assertSucceeds(parseResult: ParseResult): ParsedArgs {
+    assertThat(parseResult).isInstanceOf(ParseResult.Ok::class.java)
+    return (parseResult as ParseResult.Ok).parsedValue
+  }
+
+  private fun parseResultOk(
+      fileNames: List<String> = emptyList(),
+      formattingOptions: FormattingOptions = FormattingOptions(),
+      dryRun: Boolean = false,
+      setExitIfChanged: Boolean = false,
+      removedUnusedImports: Boolean = true,
+      stdinName: String? = null
+  ): ParseResult.Ok {
+    val returnedFormattingOptions =
+        formattingOptions.copy(removeUnusedImports = removedUnusedImports)
+    return ParseResult.Ok(
+        ParsedArgs(fileNames, returnedFormattingOptions, dryRun, setExitIfChanged, stdinName))
   }
 }