From c9d986f9a8104daf58337668f0cd37e8e5776fa1 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Thu, 15 Sep 2022 21:49:04 +0530
Subject: [PATCH] args completion for rc commands

---
 gen-go-code.py            |  2 ++
 kitty/rc/base.py          | 20 ++++++++++++++++++++
 kitty_tests/completion.py | 11 +++++++++--
 tools/completion/files.go | 33 +++++++++++++++++++++++++++++++++
 tools/completion/types.go | 11 +++++++++++
 5 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/gen-go-code.py b/gen-go-code.py
index 1539fc72e..958e9c8e4 100755
--- a/gen-go-code.py
+++ b/gen-go-code.py
@@ -51,6 +51,8 @@ def generate_completion_for_rc(name: str) -> None:
     cmd = command_for_name(name)
     if cmd.short_desc:
         print(f'{name}.Description = "{serialize_as_go_string(cmd.short_desc)}"')
+    for x in cmd.args.as_go_completion_code(name):
+        print(x)
 
 
 def generate_completions_for_kitty() -> None:
diff --git a/kitty/rc/base.py b/kitty/rc/base.py
index a4ef456cc..2aa91432d 100644
--- a/kitty/rc/base.py
+++ b/kitty/rc/base.py
@@ -197,6 +197,26 @@ class ArgsHandling:
             return 0
         return self.count
 
+    def as_go_completion_code(self, go_name: str) -> Iterator[str]:
+        from ..cli import serialize_as_go_string
+        c = self.args_count
+        if c is not None:
+            yield f'{go_name}.Stop_processing_at_arg = {c}'
+        if self.completion:
+            for k, v in self.completion.items():
+                if k == 'files':
+                    title, pats = v
+                    assert isinstance(pats, tuple)
+                    gpats = (f'"{serialize_as_go_string(x)}"' for x in pats)
+                    yield f'{go_name}.Completion_for_arg = fnmatch_completer("{serialize_as_go_string(title)}", {", ".join(gpats)})'
+                elif k == 'names':
+                    title, gen = v
+                    assert callable(gen)
+                    gpats = (f'"{serialize_as_go_string(x)}"' for x in gen())
+                    yield f'{go_name}.Completion_for_arg = names_completer("{serialize_as_go_string(title)}", {", ".join(gpats)})'
+                else:
+                    raise KeyError(f'Unknown args completion type: {k}')
+
     def as_go_code(self, cmd_name: str, field_types: Dict[str, str], handled_fields: Set[str]) -> Iterator[str]:
         c = self.args_count
         if c == 0:
diff --git a/kitty_tests/completion.py b/kitty_tests/completion.py
index a1032debf..e4cb7d839 100644
--- a/kitty_tests/completion.py
+++ b/kitty_tests/completion.py
@@ -83,9 +83,9 @@ def completion(self: TestCompletion, tdir: str):
     make_file('bin/exe1', 0o700)
     make_file('bin/exe-not1')
     make_file('exe2', 0o700)
-    make_file('exe-not2')
+    make_file('exe-not2.jpeg')
     make_file('sub/exe3', 0o700)
-    make_file('sub/exe-not3')
+    make_file('sub/exe-not3.png')
 
     add('kitty x', all_words())
     add('kitty e', all_words('exe1'))
@@ -96,6 +96,13 @@ def completion(self: TestCompletion, tdir: str):
     add('kitty ~/', all_words('~/exe3'))
     add('kitty ~/e', all_words('~/exe3'))
 
+    add('kitty @ goto-layout ', has_words('tall', 'fat'))
+    add('kitty @ goto-layout spli', all_words('splits'))
+    add('kitty @ goto-layout f f', all_words())
+    add('kitty @ set-window-logo ', all_words('exe-not2.jpeg', 'sub/exe-not3.png'))
+    add('kitty @ set-window-logo e', all_words('exe-not2.jpeg'))
+    add('kitty @ set-window-logo e e', all_words())
+
     for cmd, tests, result in zip(all_cmds, all_tests, run_tool()):
         self.current_cmd = cmd
         for test in tests:
diff --git a/tools/completion/files.go b/tools/completion/files.go
index 4b31b9c86..294582e82 100644
--- a/tools/completion/files.go
+++ b/tools/completion/files.go
@@ -78,3 +78,36 @@ func complete_executables_in_path(prefix string, paths ...string) []string {
 	}
 	return ans
 }
+
+func complete_by_fnmatch(prefix string, patterns []string) []string {
+	ans := make([]string, 0, 1024)
+	complete_files(prefix, func(completion_candidate, abspath string, d fs.DirEntry) error {
+		q := strings.ToLower(filepath.Base(abspath))
+		for _, pat := range patterns {
+			matched, err := filepath.Match(pat, q)
+			if err == nil && matched {
+				ans = append(ans, completion_candidate)
+			}
+		}
+		return nil
+	})
+	return ans
+}
+
+func fnmatch_completer(title string, patterns ...string) completion_func {
+	lpats := make([]string, 0, len(patterns))
+	for _, p := range patterns {
+		lpats = append(lpats, strings.ToLower(p))
+	}
+
+	return func(completions *Completions, word string, arg_num int) {
+		q := complete_by_fnmatch(word, lpats)
+		if len(q) > 0 {
+			mg := completions.add_match_group(title)
+			mg.IsFiles = true
+			for _, c := range q {
+				mg.add_match(c)
+			}
+		}
+	}
+}
diff --git a/tools/completion/types.go b/tools/completion/types.go
index b68e43921..f83254d92 100644
--- a/tools/completion/types.go
+++ b/tools/completion/types.go
@@ -136,3 +136,14 @@ func (self *Command) GetCompletions(argv []string) *Completions {
 	ans.Groups = non_empty_groups
 	return &ans
 }
+
+func names_completer(title string, names ...string) completion_func {
+	return func(completions *Completions, word string, arg_num int) {
+		mg := completions.add_match_group(title)
+		for _, q := range names {
+			if strings.HasPrefix(q, word) {
+				mg.add_match(q)
+			}
+		}
+	}
+}