Allow short arguments and flags to be combined (#102)

Fixes #7

This PR allows flags like `-abtrue` to mean `-a -b true`, by walking the
combined flag one character at a time and seeing if each character is a
flag that takes zero values, or a non-flag argument that takes one
value, which is set to the remaining part of the combined flag (in the
case above, `"true"`). This allows some nice shorthands, like `./mill
-wk -j10` rather than `./mill -w -k -j 10`

This is the only way I could find that allows both combining multiple
flags like `-ab`, and combining flags with values like `-btrue`. Trying
to handle more sophisticated cases like `tar tfv <file>` where `f` is
given `<file>` is fundamentally incompatible, and would need some
user-facing configuration to enable.

Combining multiple short arguments together with gflags `=` syntax, e.g.
`-ab=true`, is currently not allowed. This follows the current
limitation where we do not allow single short args to be used with `=`,
e.g. `-f=bar` is prohibited. In general, combined short arguments are
not allowed to have `=` in it. If you want the `=` to be part of the
value, you can move it into a separate token e.g. `-ab =true`

For now, this improvement can be done fully transparently and
backwards-compatibly without any need for user opt-in or configuration,
and should cover the 99% use case. There is no conflict with existing
long arguments, as those currently require two dashes `--`, and so a
multi-character token starting with a single `-` is currently
disallowed. Adding additional flexibility and configuration can come
later future

Covered by additional unit tests

Author: lihaoyi

build.sc

@@ -90,7 +90,7 @@ trait ExampleModule extends ScalaModule {
   def moduleDeps = Seq(mainargs.jvm(scala213))
 }
 
-object example {
+object example extends Module {
   object hello extends ExampleModule
   object hello2 extends ExampleModule
   object caseclass extends ExampleModule
@@ -102,4 +102,5 @@ object example {
   }
   object vararg extends ExampleModule
   object vararg2 extends ExampleModule
+  object short extends ExampleModule
 }

example/short/src/Main.scala

@@ -0,0 +1,12 @@
+package example.hello
+import mainargs.{main, arg, ParserForMethods, Flag}
+
+object Main {
+  @main
+  def bools(a: Flag, b: Boolean = false, c: Flag) = println(Seq(a.value, b, c.value))
+
+  @main
+  def strs(a: Flag, b: String) = println(Seq(a.value, b))
+
+  def main(args: Array[String]): Unit = ParserForMethods(this).runOrExit(args)
+}

mainargs/src/TokenGrouping.scala

@@ -26,24 +26,85 @@ object TokenGrouping {
       .flatMap { x => getNames(x).map(_ -> x) }
       .toMap[String, ArgSig]
 
-    lazy val keywordArgMap = makeKeywordArgMap(
-      x => x.name.map("--" + _) ++ x.shortName.map("-" + _)
-    )
+    lazy val shortArgMap: Map[Char, ArgSig] = argSigs
+      .collect{case (a, _) if !a.positional => a.shortName.map(_ -> a)}
+      .flatten
+      .toMap[Char, ArgSig]
+
+    lazy val shortFlagArgMap: Map[Char, ArgSig] = argSigs
+      .collect{case (a, r: TokensReader.Flag) if !a.positional => a.shortName.map(_ -> a)}
+      .flatten
+      .toMap[Char, ArgSig]
+
+    lazy val longKeywordArgMap = makeKeywordArgMap(_.name.map("--" + _))
+
+    def parseCombinedShortTokens(current: Map[ArgSig, Vector[String]],
+                                 head: String,
+                                 rest: List[String]) = {
+      val chars = head.drop(1)
+      var rest2 = rest
+      var i = 0
+      var currentMap = current
+      var failure: Result.Failure = null
+
+      while (i < chars.length) {
+        val c = chars(i)
+        shortFlagArgMap.get(c) match {
+          case Some(a) =>
+            // For `Flag`s in chars, we consume the char, set it to `true`, and continue
+            currentMap = Util.appendMap(currentMap, a, "")
+            i += 1
+          case None =>
+            // For other kinds of short arguments, we consume the char, set the value to
+            // the remaining characters, and exit
+            shortArgMap.get(c) match {
+              case Some(a) =>
+                if (i < chars.length - 1) {
+                  currentMap = Util.appendMap(currentMap, a, chars.drop(i + 1).stripPrefix("="))
+                } else {
+                  // If the non-flag argument is the last in the combined token, we look
+                  // ahead to grab the next token and assign it as this argument's value
+                  rest2 match {
+                    case Nil =>
+                      // If there is no next token, it is an error
+                      failure = Result.Failure.MismatchedArguments(incomplete = Some(a))
+                    case next :: remaining =>
+                      currentMap = Util.appendMap(currentMap, a, next)
+                      rest2 = remaining
+                  }
+                }
+              case None =>
+                // If we encounter a character that is neither a short flag or a
+                // short argument, it is an error
+                failure = Result.Failure.MismatchedArguments(unknown = Seq("-" + c.toString))
+            }
+            i = chars.length
+        }
 
-    lazy val longKeywordArgMap = makeKeywordArgMap(x => x.name.map("--" + _))
+      }
+
+      if (failure != null) Left(failure)
+      else Right((rest2, currentMap))
+    }
+
+    def lookupArgMap(k: String, m: Map[String, ArgSig]): Option[(ArgSig, mainargs.TokensReader[_])] = {
+      m.get(k).map(a => (a, a.reader))
+    }
 
     @tailrec def rec(
         remaining: List[String],
         current: Map[ArgSig, Vector[String]]
     ): Result[TokenGrouping[B]] = {
       remaining match {
         case head :: rest =>
+          // special handling for combined short args of the style "-xvf" or "-j10"
+          if (head.startsWith("-") && head.lift(1).exists(c => c != '-')){
+              parseCombinedShortTokens(current, head, rest) match{
+                case Left(failure) => failure
+                case Right((rest2, currentMap)) => rec(rest2, currentMap)
+            }
 
-          def lookupArgMap(k: String, m: Map[String, ArgSig]): Option[(ArgSig, mainargs.TokensReader[_])] = {
-            m.get(k).map(a => (a, a.reader))
-          }
-
-          if (head.startsWith("-") && head.exists(_ != '-')) {
+          } else if (head.startsWith("-") && head.exists(_ != '-')) {
             head.split("=", 2) match{
               case Array(first, second) =>
                 lookupArgMap(first, longKeywordArgMap) match {
@@ -54,7 +115,7 @@ object TokenGrouping {
                 }
 
               case _ =>
-                lookupArgMap(head, keywordArgMap) match {
+                lookupArgMap(head, longKeywordArgMap) match {
                   case Some((cliArg, _: TokensReader.Flag)) =>
                     rec(rest, Util.appendMap(current, cliArg, ""))
                   case Some((cliArg, _: TokensReader.Simple[_])) =>

mainargs/test/src/EqualsSyntaxTests.scala

@@ -27,18 +27,15 @@ object EqualsSyntaxTests extends TestSuite {
       ParserForMethods(Main).runOrThrow(Array("--foo=bar=qux")) ==>
         "bar=quxbar=qux false"
     }
+    test("empty") {
+      // --foo= syntax sets `foo` to an empty string
+      ParserForMethods(Main).runOrThrow(Array("--foo=")) ==>
+        " false"
+    }
     test("shortName") {
-      // -f=bar syntax doesn't work for short names
+      // -f=bar syntax does work for short names
       ParserForMethods(Main).runEither(Array("-f=bar")) ==>
-        Left(
-          """Missing argument: -f --foo <str>
-            |Unknown argument: "-f=bar"
-            |Expected Signature: run
-            |  -f --foo <str>  String to print repeatedly
-            |  --my-num <int>  How many times to print string
-            |  --bool          Example flag
-            |
-            |""".stripMargin)
+        Right("barbar false")
     }
     test("notFlags") {
       // -f=bar syntax doesn't work for flags

mainargs/test/src/FlagTests.scala

@@ -5,34 +5,132 @@ object FlagTests extends TestSuite {
 
   object Base {
     @main
-    def flaggy(a: Flag, b: Boolean, c: Flag) = a.value || b || c.value
+    def bool(a: Flag, b: Boolean, c: Flag) = Seq(a.value, b, c.value)
+    @main
+    def str(a: Flag, b: String) = Seq(a.value, b)
   }
+
   val check = new Checker(ParserForMethods(Base), allowPositional = true)
 
   val tests = Tests {
     test - check(
-      List("-b", "true"),
-      Result.Success(true)
+      List("bool", "-b", "true"),
+      Result.Success(Seq(false, true, false))
     )
     test - check(
-      List("-b", "false"),
-      Result.Success(false)
+      List("bool", "-b", "false"),
+      Result.Success(Seq(false, false, false))
     )
 
     test - check(
-      List("-a", "-b", "false"),
-      Result.Success(true)
+      List("bool", "-a", "-b", "false"),
+      Result.Success(Seq(true, false, false))
     )
 
     test - check(
-      List("-c", "-b", "false"),
-      Result.Success(true)
+      List("bool", "-c", "-b", "false"),
+      Result.Success(Seq(false, false, true))
     )
 
     test - check(
-      List("-a", "-c", "-b", "false"),
-      Result.Success(true)
+      List("bool", "-a", "-c", "-b", "false"),
+      Result.Success(Seq(true, false, true))
     )
 
+    test("combined"){
+      test - check(
+        List("bool", "-bfalse"),
+        Result.Success(List(false, false, false))
+      )
+      test - check(
+        List("bool", "-btrue"),
+        Result.Success(List(false, true, false))
+      )
+
+      test - check(
+        List("bool", "-abtrue"),
+        Result.Success(List(true, true, false))
+      )
+      test - check(
+        List("bool", "-abfalse"),
+        Result.Success(List(true, false, false))
+      )
+
+      test - check(
+        List("bool", "-a", "-btrue"),
+        Result.Success(List(true, true, false))
+      )
+
+      test - check(
+        List("bool", "-a", "-bfalse"),
+        Result.Success(List(true, false, false))
+      )
+
+      test - check(
+        List("bool", "-acbtrue"),
+        Result.Success(List(true, true, true))
+      )
+
+      test - check(
+        List("bool", "-acbfalse"),
+        Result.Success(List(true, false, true))
+      )
+
+      test - check(
+        List("bool", "-a", "-c", "-btrue"),
+        Result.Success(List(true, true, true))
+      )
+
+      test - check(
+        List("bool", "-a", "-c", "-bfalse"),
+        Result.Success(List(true, false, true))
+      )
+
+      test - check(
+        List("bool", "-a", "-btrue", "-c"),
+        Result.Success(List(true, true, true))
+      )
+
+      test - check(
+        List("bool", "-a", "-bfalse", "-c"),
+        Result.Success(List(true, false, true))
+      )
+
+      test - check(
+        List("bool", "-ba"),
+        Result.Failure.InvalidArguments(
+          List(
+            Result.ParamError.Failed(
+              new ArgSig(None, Some('b'), None, None, mainargs.TokensReader.BooleanRead, false, false),
+              Vector("a"),
+              "java.lang.IllegalArgumentException: For input string: \"a\""
+            )
+          )
+        )
+      )
+
+      test - check(
+        List("bool", "-ab"),
+        Result.Failure.MismatchedArguments(
+          Nil,
+          Nil,
+          Nil,
+          Some(
+            new ArgSig(None, Some('b'), None, None, TokensReader.BooleanRead, false, false)
+          )
+        )
+      )
+
+      test - check(List("str", "-b=value", "-a"), Result.Success(List(true, "value")))
+      test - check(List("str", "-b=", "-a"), Result.Success(List(true, "")))
+
+      test - check(List("str", "-ab=value"), Result.Success(List(true, "value")))
+
+      test - check(
+        List("str", "-bvalue", "-akey=value"),
+        Result.Failure.MismatchedArguments(Nil, List("-k"), Nil, None)
+      )
+    }
+
   }
 }

mainargs/test/src/PositionalTests.scala

@@ -44,31 +44,7 @@ object PositionalTests extends TestSuite {
     )
     test - check(
       List("-x", "true", "-y", "false", "-z", "false"),
-      Result.Failure.MismatchedArguments(
-        Vector(
-          ArgSig(
-            None,
-            Some('y'),
-            None,
-            None,
-            TokensReader.BooleanRead,
-            positional = true,
-            hidden = false
-          ),
-          ArgSig(
-            None,
-            Some('z'),
-            None,
-            None,
-            TokensReader.BooleanRead,
-            positional = false,
-            hidden = false
-          )
-        ),
-        List("-y", "false", "-z", "false"),
-        List(),
-        None
-      )
+      Result.Failure.MismatchedArguments(List(), List("-y"), List(), None)
     )
   }
 }