From 4c6e44bd5ecdc1d2cd98273b8b99fb2b2de6411e Mon Sep 17 00:00:00 2001
From: shanmu <shanmu1901@gmail.com>
Date: Sun, 30 Jun 2024 03:04:07 +0800
Subject: [PATCH] feat(clap_complete): Add support `--flag bar` and `-f bar`
 completion

---
 clap_complete/src/dynamic/completer.rs | 76 +++++++++++++++++++++++++-
 1 file changed, 74 insertions(+), 2 deletions(-)

diff --git a/clap_complete/src/dynamic/completer.rs b/clap_complete/src/dynamic/completer.rs
index ea4289d6aaa3..116bd438cddd 100644
--- a/clap_complete/src/dynamic/completer.rs
+++ b/clap_complete/src/dynamic/completer.rs
@@ -76,8 +76,75 @@ 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 = opt
+                    .map(|o| 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,
+                    })
+                    .unwrap_or(ParseState::ValueDone);
+            }
+        } else if let Some(short) = arg.to_short() {
+            let mut short = short.clone();
+            let opt = short.next_flag();
+
+            if let Some(Ok(opt)) = opt {
+                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 = opt
+                    .map(|o| match o.get_action() {
+                        clap::ArgAction::Set | clap::ArgAction::Append => {
+                            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,
+                    })
+                    .unwrap_or(ParseState::ValueDone)
+            } else {
+                state = ParseState::ValueDone
+            }
         } else {
             pos_index += 1;
             state = ParseState::ValueDone;
@@ -97,6 +164,9 @@ enum ParseState {
 
     /// Parsing a positional argument after `--`
     Pos(usize),
+
+    /// Parsing a optional flag argument
+    Opt(clap::Arg),
 }
 
 fn complete_arg(
@@ -159,6 +229,8 @@ fn complete_arg(
                     .map(|(f, help)| (format!("{}{}", dash_or_arg, f).into(), help)),
             );
         }
+    } else if let ParseState::Opt(opt) = &state {
+        completions.extend(complete_arg_value(arg.to_value(), opt, current_dir));
     }
 
     if let Some(positional) = cmd