Handle + force-push prefix in git push branch completion (#232)

Co-authored-by: Oz <oz-agent@warp.dev>
Co-authored-by: lucieleblanc <14223323+lucieleblanc@users.noreply.github.com>

Author: lucieleblanc

command-signatures/json/git.json

@@ -3991,11 +3991,12 @@
                     "generatorName": "remotes"
                 },
                 {
-                    "name": "branch",
+                    "name": "refspec",
+                    "description": "Refspec to push (supports +<branch> for force-push)",
                     "isOptional": true,
                     "generatorName": [
-                        "local_branches",
-                        "tags"
+                        "push_refspec_branches",
+                        "push_refspec_tags"
                     ]
                 }
             ]

command-signatures/src/generators/git.rs

@@ -529,15 +529,81 @@ pub fn commits_generator() -> Generator {
     )
 }
 
-pub fn local_branches_generator() -> Generator {
-    Generator::script(
-        CommandBuilder::single_command(
-            "git --no-optional-locks branch --no-color --sort=-committerdate",
-        ),
-        post_process_branches,
+fn local_branches_command() -> CommandBuilder {
+    CommandBuilder::single_command(
+        "git --no-optional-locks branch --no-color --sort=-committerdate",
     )
 }
 
+fn tags_command() -> CommandBuilder {
+    CommandBuilder::single_command("git --no-optional-locks tag --list --sort=-creatordate")
+}
+
+pub fn local_branches_generator() -> Generator {
+    Generator::script(local_branches_command(), post_process_branches)
+}
+
+fn post_process_tags(output: &str) -> GeneratorResults {
+    output
+        .lines()
+        .filter(|line| !line.is_empty())
+        .map(|line| Suggestion::with_description(line, "tag"))
+        .collect_ordered_results()
+}
+
+const FORCE_PREFIX_MARKER: &str = "__FORCE_PREFIX__";
+const FORCE_PREFIX_MARKER_LINE: &str = concat!("__FORCE_PREFIX__", "\n");
+
+/// Wraps a command to prepend a force-prefix marker if the last token starts with `+`.
+/// This handles the `git push origin +<branch>` force-push refspec syntax where
+/// the `+` prefix would otherwise prevent branch name matching.
+fn with_force_prefix_detection(
+    tokens: &[&str],
+    has_trailing_whitespace: bool,
+    cmd: CommandBuilder,
+) -> CommandBuilder {
+    if !has_trailing_whitespace && tokens.last().is_some_and(|t| t.starts_with('+')) {
+        CommandBuilder::and(
+            CommandBuilder::single_command(format!("printf '{FORCE_PREFIX_MARKER}\\n'")),
+            cmd,
+        )
+    } else {
+        cmd
+    }
+}
+
+/// Strips the force-prefix marker from generator output, returning the prefix
+/// to prepend to suggestions and the remaining output.
+fn strip_force_prefix(out: &str) -> (&str, &str) {
+    match out.strip_prefix(FORCE_PREFIX_MARKER_LINE) {
+        Some(rest) => ("+", rest),
+        None => ("", out),
+    }
+}
+
+/// Prepends a string to the `exact_string` of every suggestion in the results.
+fn prepend_to_suggestions(prefix: &str, results: &mut GeneratorResults) {
+    if !prefix.is_empty() {
+        for suggestion in &mut results.suggestions {
+            suggestion.exact_string = format!("{prefix}{}", suggestion.exact_string);
+        }
+    }
+}
+
+fn post_process_push_refspec_branches(out: &str) -> GeneratorResults {
+    let (prefix, branch_output) = strip_force_prefix(out);
+    let mut results = post_process_branches(branch_output);
+    prepend_to_suggestions(prefix, &mut results);
+    results
+}
+
+fn post_process_push_refspec_tags(out: &str) -> GeneratorResults {
+    let (prefix, tag_output) = strip_force_prefix(out);
+    let mut results = post_process_tags(tag_output);
+    prepend_to_suggestions(prefix, &mut results);
+    results
+}
+
 pub fn generator() -> CommandSignatureGenerators {
     CommandSignatureGenerators::new("git")
         .add_generator("commits", commits_generator())
@@ -664,18 +730,7 @@ pub fn generator() -> CommandSignatureGenerators {
                     .collect_unordered_results()
             }),
         )
-        .add_generator(
-            "tags",
-            Generator::script(
-                CommandBuilder::single_command("git --no-optional-locks tag --list --sort=-committerdate"),
-                |output| {
-                    output
-                        .lines()
-                        .filter(|&line| !line.is_empty()).map(|line| Suggestion::with_description(line, "tag"))
-                        .collect_ordered_results()
-                },
-            ),
-        )
+        .add_generator("tags", Generator::script(tags_command(), post_process_tags))
         .add_generator(
             "files_for_staging",
             Generator::script(
@@ -717,12 +772,41 @@ pub fn generator() -> CommandSignatureGenerators {
                     if tokens.contains(&"-r") || tokens.contains(&"--remotes") {
                         CommandBuilder::single_command("git --no-optional-locks branch -r --no-color --sort=-committerdate")
                     } else {
-                        CommandBuilder::single_command("git --no-optional-locks branch --no-color --sort=-committerdate")
+                        local_branches_command()
                     }
                 },
                 post_process_branches,
             ),
         )
+        // Generators for `git push` refspec arguments. These handle the `+` force-push prefix
+        // (e.g. `git push origin +branch`) by detecting it in the current token and prepending
+        // it to branch/tag suggestions so the completer's prefix matcher can match correctly.
+        .add_generator(
+            "push_refspec_branches",
+            Generator::command_from_tokens(
+                |tokens, has_trailing_whitespace, _| {
+                    with_force_prefix_detection(
+                        tokens,
+                        has_trailing_whitespace,
+                        local_branches_command(),
+                    )
+                },
+                post_process_push_refspec_branches,
+            ),
+        )
+        .add_generator(
+            "push_refspec_tags",
+            Generator::command_from_tokens(
+                |tokens, has_trailing_whitespace, _| {
+                    with_force_prefix_detection(
+                        tokens,
+                        has_trailing_whitespace,
+                        tags_command(),
+                    )
+                },
+                post_process_push_refspec_tags,
+            ),
+        )
         .add_alias(
             "alias",
             Alias::new(
@@ -754,7 +838,10 @@ pub fn generator() -> CommandSignatureGenerators {
 
 #[cfg(test)]
 mod tests {
-    use crate::generators::git::{post_process_branches, post_process_tracked_files};
+    use crate::generators::git::{
+        post_process_branches, post_process_push_refspec_branches, post_process_push_refspec_tags,
+        post_process_tags, post_process_tracked_files,
+    };
     use warp_completion_metadata::{
         GeneratorResults, IconType, Importance, Order, Priority, Suggestion,
     };
@@ -857,4 +944,156 @@ mod tests {
             }
         );
     }
+
+    #[test]
+    fn test_post_process_tags() {
+        let command_output = "v1.0.0\nv2.0.0\nv0.1.0";
+        assert_eq!(
+            post_process_tags(command_output),
+            GeneratorResults {
+                suggestions: vec![
+                    Suggestion {
+                        exact_string: "v1.0.0".to_owned(),
+                        display_name: None,
+                        description: Some("tag".to_owned()),
+                        priority: Priority::Default,
+                        icon: None,
+                        is_hidden: false,
+                    },
+                    Suggestion {
+                        exact_string: "v2.0.0".to_owned(),
+                        display_name: None,
+                        description: Some("tag".to_owned()),
+                        priority: Priority::Default,
+                        icon: None,
+                        is_hidden: false,
+                    },
+                    Suggestion {
+                        exact_string: "v0.1.0".to_owned(),
+                        display_name: None,
+                        description: Some("tag".to_owned()),
+                        priority: Priority::Default,
+                        icon: None,
+                        is_hidden: false,
+                    },
+                ],
+                is_ordered: true,
+            }
+        );
+    }
+
+    #[test]
+    fn test_post_process_tags_filters_empty_lines() {
+        let command_output = "v1.0.0\n\nv2.0.0\n";
+        assert_eq!(post_process_tags(command_output).suggestions.len(), 2);
+    }
+
+    #[test]
+    fn test_push_refspec_branches_without_force_prefix() {
+        // Without the __FORCE_PREFIX__ marker, results should match normal branch processing.
+        let command_output = "* main\n  feature/foo\n  develop";
+        assert_eq!(
+            post_process_push_refspec_branches(command_output),
+            post_process_branches(command_output),
+        );
+    }
+
+    #[test]
+    fn test_push_refspec_branches_with_force_prefix() {
+        // With the __FORCE_PREFIX__ marker, all branch exact_strings should be prefixed with "+".
+        let command_output = "__FORCE_PREFIX__\n* main\n  feature/foo\n  develop";
+        let results = post_process_push_refspec_branches(command_output);
+        assert_eq!(
+            results,
+            GeneratorResults {
+                suggestions: vec![
+                    Suggestion {
+                        exact_string: "+main".to_owned(),
+                        display_name: None,
+                        description: Some("Current branch".to_owned()),
+                        priority: Priority::most_important(),
+                        icon: Some(IconType::GitBranch),
+                        is_hidden: false,
+                    },
+                    Suggestion {
+                        exact_string: "+feature/foo".to_owned(),
+                        display_name: None,
+                        description: Some("Branch".to_owned()),
+                        priority: Priority::Default,
+                        icon: Some(IconType::GitBranch),
+                        is_hidden: false,
+                    },
+                    Suggestion {
+                        exact_string: "+develop".to_owned(),
+                        display_name: None,
+                        description: Some("Branch".to_owned()),
+                        priority: Priority::Default,
+                        icon: Some(IconType::GitBranch),
+                        is_hidden: false,
+                    },
+                ],
+                is_ordered: true,
+            }
+        );
+    }
+
+    #[test]
+    fn test_push_refspec_tags_without_force_prefix() {
+        let command_output = "v1.0.0\nv2.0.0";
+        let results = post_process_push_refspec_tags(command_output);
+        assert_eq!(
+            results,
+            GeneratorResults {
+                suggestions: vec![
+                    Suggestion {
+                        exact_string: "v1.0.0".to_owned(),
+                        display_name: None,
+                        description: Some("tag".to_owned()),
+                        priority: Priority::Default,
+                        icon: None,
+                        is_hidden: false,
+                    },
+                    Suggestion {
+                        exact_string: "v2.0.0".to_owned(),
+                        display_name: None,
+                        description: Some("tag".to_owned()),
+                        priority: Priority::Default,
+                        icon: None,
+                        is_hidden: false,
+                    },
+                ],
+                is_ordered: true,
+            }
+        );
+    }
+
+    #[test]
+    fn test_push_refspec_tags_with_force_prefix() {
+        let command_output = "__FORCE_PREFIX__\nv1.0.0\nv2.0.0";
+        let results = post_process_push_refspec_tags(command_output);
+        assert_eq!(
+            results,
+            GeneratorResults {
+                suggestions: vec![
+                    Suggestion {
+                        exact_string: "+v1.0.0".to_owned(),
+                        display_name: None,
+                        description: Some("tag".to_owned()),
+                        priority: Priority::Default,
+                        icon: None,
+                        is_hidden: false,
+                    },
+                    Suggestion {
+                        exact_string: "+v2.0.0".to_owned(),
+                        display_name: None,
+                        description: Some("tag".to_owned()),
+                        priority: Priority::Default,
+                        icon: None,
+                        is_hidden: false,
+                    },
+                ],
+                is_ordered: true,
+            }
+        );
+    }
 }