feat(clap_complete): Add support --flag bar and -f bar completion

clap-rs · Jul 23, 2024 · f7383f7 · f7383f7
1 parent 2f53bb3
commit f7383f7
Show file tree

Hide file tree

Showing 2 changed files with 118 additions and 66 deletions.
diff --git a/clap_complete/src/dynamic/completer.rs b/clap_complete/src/dynamic/completer.rs
@@ -76,11 +76,92 @@ pub fn complete(
         } else if arg.is_escape() {
             is_escaped = true;
             state = ParseState::ValueDone;
-        } else if let Some(_long) = arg.to_long() {
-        } else if let Some(_short) = arg.to_short() {
+        } else if let Some((flag, value)) = arg.to_long() {
+            if let Ok(flag) = flag {
+                let opt = current_cmd.get_arguments().find(|a| {
+                    let longs = a.get_long_and_visible_aliases();
+                    let is_find = longs.map(|v| {
+                        let mut iter = v.into_iter();
+                        let s = iter.find(|s| *s == flag);
+                        s.is_some()
+                    });
+                    is_find.unwrap_or(false)
+                });
+                state = match opt.map(|o| o.get_action()) {
+                    Some(clap::ArgAction::Set) | Some(clap::ArgAction::Append) => {
+                        if value.is_some() {
+                            ParseState::ValueDone
+                        } else {
+                            ParseState::Opt(opt.unwrap().clone())
+                        }
+                    }
+                    Some(clap::ArgAction::SetTrue) | Some(clap::ArgAction::SetFalse) => {
+                        ParseState::ValueDone
+                    }
+                    Some(clap::ArgAction::Count) => ParseState::ValueDone,
+                    Some(clap::ArgAction::Version) => ParseState::ValueDone,
+                    Some(clap::ArgAction::Help)
+                    | Some(clap::ArgAction::HelpLong)
+                    | Some(clap::ArgAction::HelpShort) => ParseState::ValueDone,
+                    Some(_) => ParseState::ValueDone,
+                    None => ParseState::ValueDone,
+                };
+            } else {
+                state = ParseState::ValueDone;
+            }
+        } else if let Some(mut short) = arg.to_short() {
+            let mut takes_value = false;
+            loop {
+                if let Some(Ok(opt)) = short.next_flag() {
+                    let opt = current_cmd.get_arguments().find(|a| {
+                        let shorts = a.get_short_and_visible_aliases();
+                        let is_find = shorts.map(|v| {
+                            let mut iter = v.into_iter();
+                            let c = iter.find(|c| *c == opt);
+                            c.is_some()
+                        });
+                        is_find.unwrap_or(false)
+                    });
+
+                    state = match opt.map(|o| o.get_action()) {
+                        Some(clap::ArgAction::Set) | Some(clap::ArgAction::Append) => {
+                            takes_value = true;
+                            if short.next_value_os().is_some() {
+                                ParseState::ValueDone
+                            } else {
+                                ParseState::Opt(opt.unwrap().clone())
+                            }
+                        }
+                        Some(clap::ArgAction::SetTrue) | Some(clap::ArgAction::SetFalse) => {
+                            ParseState::ValueDone
+                        }
+                        Some(clap::ArgAction::Count) => ParseState::ValueDone,
+                        Some(clap::ArgAction::Version) => ParseState::ValueDone,
+                        Some(clap::ArgAction::Help)
+                        | Some(clap::ArgAction::HelpShort)
+                        | Some(clap::ArgAction::HelpLong) => ParseState::ValueDone,
+                        Some(_) => ParseState::ValueDone,
+                        None => ParseState::ValueDone,
+                    };
+
+                    if takes_value {
+                        break;
+                    }
+                } else {
+                    state = ParseState::ValueDone;
+                    break;
+                }
+            }
         } else {
-            pos_index += 1;
-            state = ParseState::ValueDone;
+            match state {
+                ParseState::ValueDone | ParseState::Pos(_) => {
+                    pos_index += 1;
+                    state = ParseState::ValueDone;
+                }
+                ParseState::Opt(_) => {
+                    state = ParseState::ValueDone;
+                }
+            }
         }
     }
 
@@ -97,6 +178,9 @@ enum ParseState {
 
     /// Parsing a positional argument after `--`
     Pos(usize),
+
+    /// Parsing a optional flag argument
+    Opt(clap::Arg),
 }
 
 fn complete_arg(
@@ -200,6 +284,9 @@ fn complete_arg(
                 completions.extend(complete_arg_value(arg.to_value(), positional, current_dir));
             }
         }
+        ParseState::Opt(opt) => {
+            completions.extend(complete_arg_value(arg.to_value(), &opt, current_dir));
+        }
     }
     if completions.iter().any(|a| a.is_visible()) {
         completions.retain(|a| a.is_visible())

diff --git a/clap_complete/tests/testsuite/dynamic.rs b/clap_complete/tests/testsuite/dynamic.rs
@@ -324,103 +324,65 @@ fn suggest_argument_value() {
     assert_data_eq!(
         complete!(cmd, "--input [TAB]", current_dir = Some(testdir_path)),
         snapbox::str![
-            "--input
---format
---count
---help	Print help
--i
--F
--c
--h	Print help
-pos_a
-pos_b
-pos_c"
+            "a_file
+b_file
+c_dir/
+d_dir/"
         ],
     );
 
     assert_data_eq!(
         complete!(cmd, "-i [TAB]", current_dir = Some(testdir_path)),
         snapbox::str![
-            "--input
---format
---count
---help	Print help
--i
--F
--c
--h	Print help
-pos_a
-pos_b
-pos_c"
+            "a_file
+b_file
+c_dir/
+d_dir/"
         ],
     );
 
     assert_data_eq!(
         complete!(cmd, "--input a[TAB]", current_dir = Some(testdir_path)),
-        snapbox::str![""],
+        snapbox::str!["a_file"],
     );
 
     assert_data_eq!(
         complete!(cmd, "-i b[TAB]", current_dir = Some(testdir_path)),
-        snapbox::str![""],
+        snapbox::str!["b_file"],
     );
 
     assert_data_eq!(
         complete!(cmd, "--format [TAB]"),
         snapbox::str![
-            "--input
---format
---count
---help	Print help
--i
--F
--c
--h	Print help
-pos_a
-pos_b
-pos_c"
+            "json
+yaml
+toml"
         ],
     );
 
     assert_data_eq!(
         complete!(cmd, "-F [TAB]"),
         snapbox::str![
-            "--input
---format
---count
---help	Print help
--i
--F
--c
--h	Print help
-pos_a
-pos_b
-pos_c"
+            "json
+yaml
+toml"
         ],
     );
 
-    assert_data_eq!(complete!(cmd, "--format j[TAB]"), snapbox::str![""],);
+    assert_data_eq!(complete!(cmd, "--format j[TAB]"), snapbox::str!["json"],);
 
-    assert_data_eq!(complete!(cmd, "-F j[TAB]"), snapbox::str![""],);
+    assert_data_eq!(complete!(cmd, "-F j[TAB]"), snapbox::str!["json"],);
 
-    assert_data_eq!(complete!(cmd, "--format t[TAB]"), snapbox::str![""],);
+    assert_data_eq!(complete!(cmd, "--format t[TAB]"), snapbox::str!["toml"],);
 
-    assert_data_eq!(complete!(cmd, "-F t[TAB]"), snapbox::str![""],);
+    assert_data_eq!(complete!(cmd, "-F t[TAB]"), snapbox::str!["toml"],);
 
     assert_data_eq!(
         complete!(cmd, "-cccF [TAB]"),
         snapbox::str![
-            "--input
---format
---count
---help\tPrint help
--i
--F
--c
--h\tPrint help
-pos_a
-pos_b
-pos_c"
+            "json
+yaml
+toml"
         ]
     );
 
@@ -434,7 +396,10 @@ pos_c"
 -i
 -F
 -c
--h\tPrint help"
+-h\tPrint help
+pos_a
+pos_b
+pos_c"
         ]
     );
 }