Delegate based completion for kitty cmd

kitty_tests/completion.py

@@ -52,6 +52,19 @@ def all_words(*words):
     return t
 
 
+def is_delegate(num_to_remove: int = 0, command: str = ''):
+    q = {}
+    if num_to_remove:
+        q['num_to_remove'] = num_to_remove
+    if command:
+        q['command'] = command
+
+    def t(self, result):
+        d = result['delegate']
+        self.assertEqual(d, q)
+    return t
+
+
 def completion(self: TestCompletion, tdir: str):
     all_cmds = []
     all_argv = []
@@ -148,6 +161,10 @@ def completion(self: TestCompletion, tdir: str):
     make_file('editable.txt')
     add('edit-in-kitty ', has_words('editable.txt'))
 
+    add('kitty bash ', is_delegate(1, 'bash'))
+    add('kitty -1 bash ', is_delegate(2, 'bash'))
+    add('kitty -1 bash --n', is_delegate(2, 'bash'))
+
     for cmd, tests, result in zip(all_cmds, all_tests, run_tool()):
         self.current_cmd = cmd
         for test in tests:

tools/completion/kitty.go

@@ -19,6 +19,8 @@ var _ = fmt.Print
 
 func complete_kitty(completions *Completions, word string, arg_num int) {
 	if arg_num > 1 {
+		completions.Delegate.NumToRemove = completions.current_cmd.index_of_first_arg + 1 // +1 because the first word is not present in all_words
+		completions.Delegate.Command = completions.all_words[completions.current_cmd.index_of_first_arg]
 		return
 	}
 	exes := complete_executables_in_path(word)

tools/completion/parse-args.go

@@ -146,8 +146,11 @@ func default_parse_args(cmd *Command, words []string, completions *Completions)
 		completions.current_word_idx = i
 		completions.current_word_idx_in_parent++
 		is_last_word := i == len(words)-1
-		if expecting_arg_for == nil && !strings.HasPrefix(word, "-") {
+		if only_args_allowed || (expecting_arg_for == nil && !strings.HasPrefix(word, "-")) {
 			arg_num++
+			if arg_num == 1 {
+				cmd.index_of_first_arg = completions.current_word_idx
+			}
 		}
 		if is_last_word {
 			complete_word(word, completions, only_args_allowed, expecting_arg_for, arg_num)

tools/completion/types.go

@@ -72,8 +72,14 @@ func (self *MatchGroup) longest_common_prefix() string {
 	}, true)
 }
 
+type Delegate struct {
+	NumToRemove int    `json:"num_to_remove,omitempty"`
+	Command     string `json:"command,omitempty"`
+}
+
 type Completions struct {
 	Groups   []*MatchGroup `json:"groups,omitempty"`
+	Delegate Delegate      `json:"delegate,omitempty"`
 
 	current_cmd                *Command
 	all_words                  []string // all words passed to parse_args()
@@ -126,6 +132,10 @@ type Command struct {
 	Subcommand_must_be_first        bool
 
 	Parse_args func(*Command, []string, *Completions)
+
+	// index in Completions.all_words of the first non-option argument to this command.
+	// A value of zero means no arg was found while parsing.
+	index_of_first_arg int
 }
 
 func (self *Command) clone_options_from(other *Command) {

tools/completion/zsh.go

@@ -75,6 +75,14 @@ func fmt_desc(word, desc string, max_word_len int, f *markup.Context, screen_wid
 func serialize(completions *Completions, f *markup.Context, screen_width int) ([]byte, error) {
 	cmd := strings.Builder{}
 	output := strings.Builder{}
+	if completions.Delegate.NumToRemove > 0 {
+		for i := 0; i < completions.Delegate.NumToRemove; i++ {
+			fmt.Fprintln(&output, "shift words")
+			fmt.Fprintln(&output, "(( CURRENT-- ))")
+		}
+		fmt.Fprintln(&output, "_normal -p ", utils.QuoteStringForSH(completions.Delegate.Command))
+		return []byte(output.String()), nil
+	}
 	for _, mg := range completions.Groups {
 		cmd.WriteString("compadd -U -J ")
 		cmd.WriteString(utils.QuoteStringForSH(mg.Title))