From 936d89e69471faee14efbd28b37430ea0db0aff3 Mon Sep 17 00:00:00 2001
From: Li Haoyi <haoyi.sg@gmail.com>
Date: Fri, 13 Sep 2024 11:08:50 +0800
Subject: [PATCH] Fix error reporting of incomplete short args (#162)

This was accidentally broken in
https://github.com/com-lihaoyi/mainargs/pull/102 and
https://github.com/com-lihaoyi/mainargs/pull/112, and wasn't covered by
tests. Noticed when trying to update Ammonite to the latest version of
MainArgs in https://github.com/com-lihaoyi/Ammonite/pull/1549

Restored the special casing for tracking/handling incomplete arguments
and added some unit test cases
---
 mainargs/src/TokenGrouping.scala  | 10 ++++---
 mainargs/test/src/CoreTests.scala | 45 ++++++++++++++++++++++++++-----
 mainargs/test/src/FlagTests.scala |  6 ++---
 3 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/mainargs/src/TokenGrouping.scala b/mainargs/src/TokenGrouping.scala
index 8ff345a..faa3d96 100644
--- a/mainargs/src/TokenGrouping.scala
+++ b/mainargs/src/TokenGrouping.scala
@@ -58,6 +58,7 @@ object TokenGrouping {
       var i = 0
       var currentMap = current
       var failure = false
+      var incomplete: Option[ArgSig] = None
 
       while (i < chars.length) {
         val c = chars(i)
@@ -79,6 +80,7 @@ object TokenGrouping {
                   rest2 match {
                     case Nil =>
                       // If there is no next token, it is an error
+                      incomplete = Some(a)
                       failure = true
                     case next :: remaining =>
                       currentMap = Util.appendMap(currentMap, a, next)
@@ -95,7 +97,7 @@ object TokenGrouping {
 
       }
 
-      if (failure) None else Some((rest2, currentMap))
+      if (failure) Left(incomplete) else Right((rest2, currentMap))
     }
 
     def lookupArgMap(k: String, m: Map[String, ArgSig]): Option[(ArgSig, mainargs.TokensReader[_])] = {
@@ -111,8 +113,10 @@ object TokenGrouping {
           // 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 None => complete(remaining, current)
-                case Some((rest2, currentMap)) => rec(rest2, currentMap)
+                case Left(Some(incompleteArg)) =>
+                  Result.Failure.MismatchedArguments(Nil, Nil, Nil, incomplete = Some(incompleteArg))
+                case Left(None) => complete(remaining, current)
+                case Right((rest2, currentMap)) => rec(rest2, currentMap)
             }
 
           } else if (head.startsWith("-") && head.exists(_ != '-')) {
diff --git a/mainargs/test/src/CoreTests.scala b/mainargs/test/src/CoreTests.scala
index df6eaeb..803e1f1 100644
--- a/mainargs/test/src/CoreTests.scala
+++ b/mainargs/test/src/CoreTests.scala
@@ -15,7 +15,11 @@ object CoreBase {
       i: Int,
       @arg(doc = "Pass in a custom `s` to override it")
       s: String = "lols"
-  ) = s * i
+  )
+  = s * i
+  @main
+  def baz(arg: Int) = arg
+
   @main
   def ex() = throw MyException
 
@@ -44,7 +48,10 @@ class CoreTests(allowPositional: Boolean) extends TestSuite {
           |  qux
           |  Qux is a function that does stuff
           |    -i <int>
-          |    -s <str>  Pass in a custom `s` to override it
+          |    -s <str>     Pass in a custom `s` to override it
+          |
+          |  baz
+          |    --arg <int>
           |
           |  ex
           |""".stripMargin
@@ -56,27 +63,31 @@ class CoreTests(allowPositional: Boolean) extends TestSuite {
 
       assert(
         names ==
-          List("foo", "bar", "qux", "ex")
+          List("foo", "bar", "qux", "baz", "ex")
       )
       val evaledArgs = check.mains.value.map(_.flattenedArgSigs.map {
-        case (ArgSig(name, s, docs, None, parser, _, _), _) => (s, docs, None, parser)
+        case (ArgSig(name, s, docs, None, parser, _, _), _) => (name, s, docs, None, parser)
         case (ArgSig(name, s, docs, Some(default), parser, _, _), _) =>
-          (s, docs, Some(default(CoreBase)), parser)
+          (name, s, docs, Some(default(CoreBase)), parser)
       })
 
       assert(
         evaledArgs == List(
           List(),
-          List((Some('i'), None, None, TokensReader.IntRead)),
+          List((None, Some('i'), None, None, TokensReader.IntRead)),
           List(
-            (Some('i'), None, None, TokensReader.IntRead),
+            (None, Some('i'), None, None, TokensReader.IntRead),
             (
+              None,
               Some('s'),
               Some("Pass in a custom `s` to override it"),
               Some("lols"),
               TokensReader.StringRead
             )
           ),
+          List(
+            (Some("arg"), None, None, None, TokensReader.IntRead),
+          ),
           List()
         )
       )
@@ -127,6 +138,26 @@ class CoreTests(allowPositional: Boolean) extends TestSuite {
                 None
               ) =>
         }
+        test("incomplete") {
+          // Make sure both long args and short args properly report
+          // incomplete arguments as distinct from other mismatches
+          test - assertMatch(check.parseInvoke(List("qux", "-s"))) {
+            case Result.Failure.MismatchedArguments(
+            Nil,
+            Nil,
+            Nil,
+            Some(_)
+            ) =>
+          }
+          test - assertMatch(check.parseInvoke(List("baz", "--arg"))) {
+            case Result.Failure.MismatchedArguments(
+            Nil,
+            Nil,
+            Nil,
+            Some(_)
+            ) =>
+          }
+        }
       }
 
       test("tooManyParams") - check(
diff --git a/mainargs/test/src/FlagTests.scala b/mainargs/test/src/FlagTests.scala
index 8a6bef2..459a2ce 100644
--- a/mainargs/test/src/FlagTests.scala
+++ b/mainargs/test/src/FlagTests.scala
@@ -112,10 +112,10 @@ object FlagTests extends TestSuite {
       test - check(
         List("bool", "-ab"),
         Result.Failure.MismatchedArguments(
-          Vector(new ArgSig(None, Some('b'), None, None, TokensReader.BooleanRead, false, false)),
-          List("-ab"),
           Nil,
-          None
+          Nil,
+          Nil,
+          Some(new ArgSig(None, Some('b'), None, None, TokensReader.BooleanRead, false, false))
         )
       )