From f0bd47506f0bca749c41c3af6f3ca63589992607 Mon Sep 17 00:00:00 2001
From: shannmu <shanmu1901@gmail.com>
Date: Fri, 26 Jul 2024 20:08:16 +0800
Subject: [PATCH] feat(clap_complete): Support multi-values of positional
 argument with `num_arg(N)`

---
 clap_complete/src/dynamic/completer.rs   | 68 +++++++++++++++++++++---
 clap_complete/tests/testsuite/dynamic.rs | 33 +++++++-----
 2 files changed, 79 insertions(+), 22 deletions(-)

diff --git a/clap_complete/src/dynamic/completer.rs b/clap_complete/src/dynamic/completer.rs
index 1c711e810fd..7a9015399a8 100644
--- a/clap_complete/src/dynamic/completer.rs
+++ b/clap_complete/src/dynamic/completer.rs
@@ -1,3 +1,4 @@
+use core::num;
 use std::ffi::OsStr;
 use std::ffi::OsString;
 
@@ -71,8 +72,7 @@ pub fn complete(
         }
 
         if is_escaped {
-            pos_index += 1;
-            state = ParseState::Pos(pos_index);
+            (state, pos_index) = parse_positional(current_cmd, pos_index, is_escaped, state);
         } else if arg.is_escape() {
             is_escaped = true;
             state = ParseState::ValueDone;
@@ -124,10 +124,11 @@ pub fn complete(
             }
         } else {
             match state {
-                ParseState::ValueDone | ParseState::Pos(_) => {
-                    pos_index += 1;
-                    state = ParseState::ValueDone;
+                ParseState::ValueDone | ParseState::Pos(..) => {
+                    (state, pos_index) =
+                        parse_positional(current_cmd, pos_index, is_escaped, state);
                 }
+
                 ParseState::Opt((ref opt, count)) => match opt.get_num_args() {
                     Some(range) => {
                         let max = range.max_values();
@@ -156,8 +157,8 @@ enum ParseState<'a> {
     /// Parsing a value done, there is no state to record.
     ValueDone,
 
-    /// Parsing a positional argument after `--`
-    Pos(usize),
+    /// Parsing a positional argument after `--`. Pos(pos_index, takes_num_args)
+    Pos((usize, usize)),
 
     /// Parsing a optional flag argument
     Opt((&'a clap::Arg, usize)),
@@ -300,7 +301,7 @@ fn complete_arg(
                 completions.extend(complete_subcommand(value, cmd));
             }
         }
-        ParseState::Pos(_) => {
+        ParseState::Pos(..) => {
             if let Some(positional) = cmd
                 .get_positionals()
                 .find(|p| p.get_index() == Some(pos_index))
@@ -603,6 +604,57 @@ fn parse_shortflags<'c, 's>(
     (leading_flags, takes_value_opt, short)
 }
 
+/// Parse the positional arguments. Return the new state and the new positional index.
+fn parse_positional<'a>(
+    cmd: &clap::Command,
+    pos_index: usize,
+    is_escaped: bool,
+    state: ParseState<'a>,
+) -> (ParseState<'a>, usize) {
+    let pos_arg = cmd
+        .get_positionals()
+        .find(|p| p.get_index() == Some(pos_index));
+    let num_args = pos_arg
+        .and_then(|a| a.get_num_args().and_then(|r| Some(r.max_values())))
+        .unwrap_or(1);
+
+    let update_state_with_new_positional = |pos_index| -> (ParseState<'a>, usize) {
+        if num_args > 1 {
+            (ParseState::Pos((pos_index, 1)), pos_index)
+        } else {
+            if is_escaped {
+                (ParseState::Pos((pos_index, 1)), pos_index + 1)
+            } else {
+                (ParseState::ValueDone, pos_index + 1)
+            }
+        }
+    };
+    match state {
+        ParseState::ValueDone => {
+            update_state_with_new_positional(pos_index)
+        },
+        ParseState::Pos((prev_pos_index, num_arg)) => {
+            if prev_pos_index == pos_index {
+                if num_arg + 1 < num_args {
+                    (ParseState::Pos((pos_index, num_arg + 1)), pos_index)
+                } else {
+                    if is_escaped {
+                        (ParseState::Pos((pos_index, 1)), pos_index + 1)
+                    } else {
+                        (ParseState::ValueDone, pos_index + 1)
+                    }
+                }
+            } else {
+                update_state_with_new_positional(pos_index)
+            }
+        }
+        ParseState::Opt(..) => unreachable!(
+            "This branch won't be hit, 
+            because ParseState::Opt should not be seen as a positional argument and passed to this function."
+        ),
+    }
+}
+
 /// A completion candidate definition
 ///
 /// This makes it easier to add more fields to completion candidate,
diff --git a/clap_complete/tests/testsuite/dynamic.rs b/clap_complete/tests/testsuite/dynamic.rs
index a843b78b320..0c0fe22b720 100644
--- a/clap_complete/tests/testsuite/dynamic.rs
+++ b/clap_complete/tests/testsuite/dynamic.rs
@@ -614,20 +614,18 @@ fn suggest_multi_positional() {
     assert_data_eq!(
         complete!(cmd, "pos_1 pos_a [TAB]"),
         snapbox::str![
-            "--format
---help\tPrint help
--F
--h\tPrint help"
+            "pos_a
+pos_b
+pos_c"
         ]
     );
 
     assert_data_eq!(
         complete!(cmd, "pos_1 pos_a pos_b [TAB]"),
         snapbox::str![
-            "--format
---help\tPrint help
--F
--h\tPrint help"
+            "pos_a
+pos_b
+pos_c"
         ]
     );
 
@@ -647,10 +645,9 @@ pos_c"
     assert_data_eq!(
         complete!(cmd, "--format json pos_1 pos_a [TAB]"),
         snapbox::str![
-            "--format
---help\tPrint help
--F
--h\tPrint help"
+            "pos_a
+pos_b
+pos_c"
         ]
     );
 
@@ -666,12 +663,20 @@ pos_c"
 
     assert_data_eq!(
         complete!(cmd, "--format json -- pos_1 pos_a [TAB]"),
-        snapbox::str![""]
+        snapbox::str![
+            "pos_a
+pos_b
+pos_c"
+        ]
     );
 
     assert_data_eq!(
         complete!(cmd, "--format json -- pos_1 pos_a pos_b [TAB]"),
-        snapbox::str![""]
+        snapbox::str![
+            "pos_a
+pos_b
+pos_c"
+        ]
     );
 
     assert_data_eq!(