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

clap-rs · Jul 12, 2024 · 97fde8f · 97fde8f
1 parent a116815
commit 97fde8f
Show file tree

Hide file tree

Showing 2 changed files with 116 additions and 32 deletions.
diff --git a/clap_complete/src/dynamic/completer.rs b/clap_complete/src/dynamic/completer.rs
@@ -76,8 +76,90 @@ 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 = if let Some(o) = opt {
+                    match o.get_action() {
+                        clap::ArgAction::Set | clap::ArgAction::Append => {
+                            if value.is_some() {
+                                ParseState::ValueDone
+                            } else {
+                                ParseState::Opt(o.clone())
+                            }
+                        }
+                        clap::ArgAction::SetTrue | clap::ArgAction::SetFalse => {
+                            ParseState::ValueDone
+                        }
+                        clap::ArgAction::Count => ParseState::ValueDone,
+                        clap::ArgAction::Version => ParseState::ValueDone,
+                        clap::ArgAction::Help
+                        | clap::ArgAction::HelpLong
+                        | clap::ArgAction::HelpShort => ParseState::ValueDone,
+                        _ => ParseState::ValueDone,
+                    }
+                } else {
+                    ParseState::ValueDone
+                };
+            } else {
+                state = ParseState::ValueDone;
+            }
+        } else if let Some(short) = arg.to_short() {
+            let mut short = short.clone();
+
+            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)
+                    });
+                    let mut flag = false;
+                    state = if let Some(o) = opt {
+                        match o.get_action() {
+                            clap::ArgAction::Set | clap::ArgAction::Append => {
+                                flag = true;
+                                if short.next_value_os().is_some() {
+                                    ParseState::ValueDone
+                                } else {
+                                    ParseState::Opt(o.clone())
+                                }
+                            }
+                            clap::ArgAction::SetTrue | clap::ArgAction::SetFalse => {
+                                ParseState::ValueDone
+                            }
+                            clap::ArgAction::Count => ParseState::ValueDone,
+                            clap::ArgAction::Version => ParseState::ValueDone,
+                            clap::ArgAction::Help
+                            | clap::ArgAction::HelpShort
+                            | clap::ArgAction::HelpLong => ParseState::ValueDone,
+                            _ => ParseState::ValueDone,
+                        }
+                    } else {
+                        ParseState::ValueDone
+                    };
+
+                    if flag {
+                        break;
+                    }
+                } else {
+                    state = ParseState::ValueDone;
+                    break;
+                }
+            }
         } else {
             pos_index += 1;
             state = ParseState::ValueDone;
@@ -97,6 +179,9 @@ enum ParseState {
 
     /// Parsing a positional argument after `--`
     Pos(usize),
+
+    /// Parsing a optional flag argument
+    Opt(clap::Arg),
 }
 
 fn complete_arg(
@@ -190,6 +275,9 @@ fn complete_arg(
                 completions.extend(complete_subcommand(value, cmd));
             }
         }
+        ParseState::Opt(opt) => {
+            completions.extend(complete_arg_value(arg.to_value(), &opt, current_dir));
+        }
     }
 
     Ok(completions)

diff --git a/clap_complete/tests/testsuite/dynamic.rs b/clap_complete/tests/testsuite/dynamic.rs
@@ -148,6 +148,12 @@ fn suggest_argument_value() {
                 .short('F')
                 .value_parser(["json", "yaml", "toml"]),
         )
+        .arg(
+            clap::Arg::new("count")
+                .long("count")
+                .short('c')
+                .action(clap::ArgAction::Count),
+        )
         .args_conflicts_with_subcommands(true);
 
     let testdir = snapbox::dir::DirRoot::mutable_temp().unwrap();
@@ -161,68 +167,58 @@ fn suggest_argument_value() {
     assert_data_eq!(
         complete!(cmd, "--input [TAB]", current_dir = Some(testdir_path)),
         snapbox::str![
-            "--input
---format
---help	Print help
--i
--F
--h	Print help"
+            "a_file
+b_file
+c_dir/
+d_dir/"
         ],
     );
 
     assert_data_eq!(
         complete!(cmd, "-i [TAB]", current_dir = Some(testdir_path)),
         snapbox::str![
-            "--input
---format
---help	Print help
--i
--F
--h	Print help"
+            "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
---help	Print help
--i
--F
--h	Print help"
+            "json
+yaml
+toml"
         ],
     );
 
     assert_data_eq!(
         complete!(cmd, "-F [TAB]"),
         snapbox::str![
-            "--input
---format
---help	Print help
--i
--F
--h	Print help"
+            "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"],);
 }
 
 fn complete(cmd: &mut Command, args: impl AsRef<str>, current_dir: Option<&Path>) -> String {