From 928a4db81792318d9a2ebbe9073b0e14b5bfeadf Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 19 Sep 2022 13:31:32 +0530
Subject: [PATCH] Implement delegate based completion fof the ssh and
 hyperlinked_grep kittens

---
 gen-go-code.py                   | 10 +++++++---
 kittens/hyperlinked_grep/main.py |  3 +++
 kittens/runner.py                |  8 ++++++++
 kittens/ssh/main.py              |  6 +++---
 kitty_tests/completion.py        |  1 +
 tools/completion/kitty.go        |  7 +++++++
 tools/completion/parse-args.go   |  4 ++++
 7 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/gen-go-code.py b/gen-go-code.py
index 70aa39cdc..0f1a5c5bb 100755
--- a/gen-go-code.py
+++ b/gen-go-code.py
@@ -59,10 +59,14 @@ def generate_completion_for_rc(name: str) -> None:
 
 
 def generate_kittens_completion() -> None:
-    from kittens.runner import all_kitten_names, get_kitten_cli_docs
-    for kitten in all_kitten_names():
+    from kittens.runner import all_kitten_names, get_kitten_cli_docs, get_kitten_wrapper_of
+    for kitten in sorted(all_kitten_names()):
         kn = 'kitten_' + kitten
         print(f'{kn} := plus_kitten.add_command("{kitten}", "Kittens")')
+        wof = get_kitten_wrapper_of(kitten)
+        if wof:
+            print(f'{kn}.Parse_args = completion_for_wrapper("{serialize_as_go_string(wof)}")')
+            continue
         kcd = get_kitten_cli_docs(kitten)
         if kcd:
             ospec = kcd['options']
@@ -127,7 +131,7 @@ def generate_completions_for_kitty() -> None:
     # kitten @
     print('at := k.add_command("@", "Remote control")')
     print('at.Description = "Control kitty using commands"')
-    for go_name in all_command_names():
+    for go_name in sorted(all_command_names()):
         name = go_name.replace('_', '-')
         print(f'{go_name} := at.add_command("{name}", "")')
         generate_completion_for_rc(go_name)
diff --git a/kittens/hyperlinked_grep/main.py b/kittens/hyperlinked_grep/main.py
index fb28fafd0..c2c4fa801 100755
--- a/kittens/hyperlinked_grep/main.py
+++ b/kittens/hyperlinked_grep/main.py
@@ -112,3 +112,6 @@ def main() -> None:
 
 if __name__ == '__main__':
     main()
+elif __name__ == '__wrapper_of__':
+    cd = sys.cli_docs  # type: ignore
+    cd['wrapper_of'] = 'rg'
diff --git a/kittens/runner.py b/kittens/runner.py
index 96d271665..b99106c14 100644
--- a/kittens/runner.py
+++ b/kittens/runner.py
@@ -155,6 +155,14 @@ def get_kitten_cli_docs(kitten: str) -> Any:
         return ans
 
 
+def get_kitten_wrapper_of(kitten: str) -> str:
+    setattr(sys, 'cli_docs', {})
+    run_kitten(kitten, run_name='__wrapper_of__')
+    ans = getattr(sys, 'cli_docs')
+    delattr(sys, 'cli_docs')
+    return ans.get('wrapper_of') or ''
+
+
 def get_kitten_completer(kitten: str) -> Any:
     run_kitten(kitten, run_name='__completer__')
     ans = getattr(sys, 'kitten_completer', None)
diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index f4a8dd3f1..6be788b1a 100644
--- a/kittens/ssh/main.py
+++ b/kittens/ssh/main.py
@@ -754,9 +754,9 @@ def main(args: List[str]) -> None:
 
 if __name__ == '__main__':
     main(sys.argv)
-elif __name__ == '__completer__':
-    from .completion import complete
-    setattr(sys, 'kitten_completer', complete)
+elif __name__ == '__wrapper_of__':
+    cd = sys.cli_docs  # type: ignore
+    cd['wrapper_of'] = 'ssh'
 elif __name__ == '__conf__':
     from .options.definition import definition
     sys.options_definition = definition  # type: ignore
diff --git a/kitty_tests/completion.py b/kitty_tests/completion.py
index d9dc08c39..d643f177d 100644
--- a/kitty_tests/completion.py
+++ b/kitty_tests/completion.py
@@ -156,6 +156,7 @@ def completion(self: TestCompletion, tdir: str):
     add('kitty + kitten diff ', has_words('exe-not2.jpeg'))
     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('clone-in-kitty --ty', has_words('--type'))
     make_file('editable.txt')
diff --git a/tools/completion/kitty.go b/tools/completion/kitty.go
index 51231e600..a45ed1c22 100644
--- a/tools/completion/kitty.go
+++ b/tools/completion/kitty.go
@@ -122,3 +122,10 @@ func complete_themes(completions *Completions, word string, arg_num int) {
 		}
 	}
 }
+
+func completion_for_wrapper(wrapped_cmd string) func(*Command, []string, *Completions) {
+	return func(cmd *Command, args []string, completions *Completions) {
+		completions.Delegate.NumToRemove = completions.current_word_idx + 1
+		completions.Delegate.Command = wrapped_cmd
+	}
+}
diff --git a/tools/completion/parse-args.go b/tools/completion/parse-args.go
index 39e7fc010..eafe0ebb1 100644
--- a/tools/completion/parse-args.go
+++ b/tools/completion/parse-args.go
@@ -183,6 +183,10 @@ func default_parse_args(cmd *Command, words []string, completions *Completions)
 				arg_num = 0
 				completions.current_word_idx_in_parent = 0
 				only_args_allowed = false
+				if cmd.Parse_args != nil {
+					cmd.Parse_args(cmd, words[i+1:], completions)
+					return
+				}
 			} else if cmd.Stop_processing_at_arg > 0 && arg_num >= cmd.Stop_processing_at_arg {
 				return
 			} else {