From 3bd4fd999acf718f29917494ac069efb3ea43d26 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 26 Sep 2022 11:54:57 +0530
Subject: [PATCH] All completion tests pass again

---
 gen-go-code.py                     | 1 +
 kitty_tests/completion.py          | 3 ++-
 tools/cli/command.go               | 2 ++
 tools/cli/completion-parse-args.go | 2 +-
 tools/cli/completion.go            | 2 +-
 5 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/gen-go-code.py b/gen-go-code.py
index 423f86b69..e317b5772 100755
--- a/gen-go-code.py
+++ b/gen-go-code.py
@@ -58,6 +58,7 @@ def generate_kittens_completion() -> None:
         wof = get_kitten_wrapper_of(kitten)
         if wof:
             print(f'{kn}.ArgCompleter = cli.CompletionForWrapper("{serialize_as_go_string(wof)}")')
+            print(f'{kn}.OnlyArgsAllowed = true')
             continue
         kcd = get_kitten_cli_docs(kitten)
         if kcd:
diff --git a/kitty_tests/completion.py b/kitty_tests/completion.py
index b97a4b530..ad184286c 100644
--- a/kitty_tests/completion.py
+++ b/kitty_tests/completion.py
@@ -61,7 +61,7 @@ def is_delegate(num_to_remove: int = 0, command: str = ''):
 
     def t(self, result):
         d = result['delegate']
-        self.assertEqual(d, q)
+        self.assertEqual(d, q, f'Command line: {self.current_cmd!r}')
     return t
 
 
@@ -157,6 +157,7 @@ def completion(self: TestCompletion, tdir: str):
     add('kitty + kitten themes --', has_words('--cache-age'))
     add('kitty + kitten themes D', has_words('Default'))
     add('kitty + kitten hyperlinked_grep ', is_delegate(3, 'rg'))
+    add('kitty +kitten hyperlinked_grep ', is_delegate(2, 'rg'))
 
     add('clone-in-kitty --ty', has_words('--type'))
     make_file('editable.txt')
diff --git a/tools/cli/command.go b/tools/cli/command.go
index a0d99f5ef..d19815ee9 100644
--- a/tools/cli/command.go
+++ b/tools/cli/command.go
@@ -31,6 +31,8 @@ type Command struct {
 	ArgCompleter CompletionFunc
 	// Stop completion processing at this arg num
 	StopCompletingAtArg int
+	// Consider all args as non-options args
+	OnlyArgsAllowed bool
 	// Specialised arg aprsing
 	ParseArgsForCompletion func(cmd *Command, args []string, completions *Completions)
 
diff --git a/tools/cli/completion-parse-args.go b/tools/cli/completion-parse-args.go
index 0d79e4ad3..8119b3e97 100644
--- a/tools/cli/completion-parse-args.go
+++ b/tools/cli/completion-parse-args.go
@@ -187,7 +187,7 @@ func completion_parse_args(cmd *Command, words []string, completions *Completion
 				cmd = sc
 				arg_num = 0
 				completions.CurrentWordIdxInParent = 0
-				only_args_allowed = false
+				only_args_allowed = cmd.OnlyArgsAllowed
 				if cmd.ParseArgsForCompletion != nil {
 					cmd.ParseArgsForCompletion(cmd, words[i+1:], completions)
 					return
diff --git a/tools/cli/completion.go b/tools/cli/completion.go
index 93e5ac3ba..9dc269b96 100644
--- a/tools/cli/completion.go
+++ b/tools/cli/completion.go
@@ -155,7 +155,7 @@ func ChainCompleters(completers ...CompletionFunc) CompletionFunc {
 
 func CompletionForWrapper(wrapped_cmd string) func(completions *Completions, word string, arg_num int) {
 	return func(completions *Completions, word string, arg_num int) {
-		completions.Delegate.NumToRemove = completions.CurrentWordIdx + 1
+		completions.Delegate.NumToRemove = completions.CurrentWordIdx
 		completions.Delegate.Command = wrapped_cmd
 	}
 }