Start work on implementing edit-in-kitty in kitty-tool

gen-go-code.py

@@ -8,7 +8,7 @@ import subprocess
 import sys
 from contextlib import contextmanager, suppress
 from functools import lru_cache
-from typing import Dict, Iterator, List, Set, Tuple, Union
+from typing import Dict, Iterator, List, Set, Tuple, Union, Sequence
 
 import kitty.constants as kc
 from kittens.tui.operations import Mode
@@ -74,14 +74,21 @@ def generate_kittens_completion() -> None:
             print(f'{kn}.HelpText = ""')
 
 
-def completion_for_launch_wrappers(*names: str) -> None:
+@lru_cache
+def clone_safe_launch_opts() -> Sequence[GoOption]:
     from kitty.launch import clone_safe_opts, options_spec
-    opts = tuple(go_options_for_seq(parse_option_spec(options_spec())[0]))
+    ans = []
     allowed = clone_safe_opts()
-    for o in opts:
+    for o in go_options_for_seq(parse_option_spec(options_spec())[0]):
         if o.obj_dict['name'] in allowed:
-            for name in names:
-                print(o.as_option(name))
+            ans.append(o)
+    return tuple(ans)
+
+
+def completion_for_launch_wrappers(*names: str) -> None:
+    for o in clone_safe_launch_opts():
+        for name in names:
+            print(o.as_option(name))
 
 
 def generate_completions_for_kitty() -> None:
@@ -391,6 +398,14 @@ def update_completion() -> None:
         with replace_if_needed('tools/cmd/completion/kitty_generated.go') as f:
             sys.stdout = f
             generate_completions_for_kitty()
+        with replace_if_needed('tools/cmd/edit_in_kitty/launch_generated.go') as f:
+            sys.stdout = f
+            print('package edit_in_kitty')
+            print('import "kitty/tools/cli"')
+            print('func AddCloneSafeOpts(cmd *cli.Command) {')
+            completion_for_launch_wrappers('cmd')
+            print(''.join(CompletionSpec.from_string('type:file mime:text/* group:"Text files"').as_go_code('cmd.ArgCompleter', ' = ')))
+            print('}')
     finally:
         sys.stdout = orig
 

kitty/launch.py

@@ -672,9 +672,10 @@ class EditCmd:
         self.file_data = b''
         self.file_inode = -1, -1
         self.file_size = -1
+        self.version = 0
         self.source_window_id = self.editor_window_id = -1
         self.abort_signaled = ''
-        simple = 'file_inode', 'file_data', 'abort_signaled'
+        simple = 'file_inode', 'file_data', 'abort_signaled', 'version'
         for k, v in parse_message(msg, simple):
             if k == 'file_inode':
                 q = map(int, v.split(':'))
@@ -685,10 +686,14 @@ class EditCmd:
             elif k == 'file_data':
                 import base64
                 self.file_data = base64.standard_b64decode(v)
+            elif k == 'version':
+                self.version = int(v)
             else:
                 setattr(self, k, v)
         if self.abort_signaled:
             return
+        if self.version > 0:
+            raise ValueError(f'Unsupported version received in edit protocol: {self.version}')
         self.opts, extra_args = parse_opts_for_clone(['--type=overlay'] + self.args)
         self.file_spec = extra_args.pop()
         self.line_number = 0

tools/cmd/edit_in_kitty/main.go

@@ -0,0 +1,253 @@
+// License: GPLv3 Copyright: 2022, Kovid Goyal, <kovid at kovidgoyal.net>
+
+package edit_in_kitty
+
+import (
+	"encoding/base64"
+	"fmt"
+	"io"
+	"io/fs"
+	"os"
+	"strconv"
+	"strings"
+
+	"golang.org/x/sys/unix"
+
+	"kitty/tools/cli"
+	"kitty/tools/tui"
+	"kitty/tools/tui/loop"
+	"kitty/tools/utils"
+	"kitty/tools/utils/humanize"
+)
+
+var _ = fmt.Print
+
+func encode(x string) string {
+	return base64.StdEncoding.EncodeToString(utils.UnsafeStringToBytes(x))
+}
+
+type OnDataCallback = func(data_type string, data []byte) error
+
+func edit_loop(data_to_send string, kill_if_signaled bool, on_data OnDataCallback) (err error) {
+	lp, err := loop.New(loop.NoAlternateScreen, loop.NoRestoreColors, loop.NoMouseTracking)
+	if err != nil {
+		return
+	}
+	current_text := strings.Builder{}
+	data := strings.Builder{}
+	data.Grow(4096)
+	started := false
+	canceled := false
+	update_type := ""
+
+	handle_line := func(line string) error {
+		if canceled {
+			return nil
+		}
+		if started {
+			if update_type == "" {
+				update_type = line
+			} else {
+				if line == "KITTY_DATA_END" {
+					lp.QueueWriteString(update_type + "\r\n")
+					if update_type == "DONE" {
+						lp.Quit(0)
+						return nil
+					}
+					b, err := base64.StdEncoding.DecodeString(data.String())
+					data.Reset()
+					data.Grow(4096)
+					started = false
+					if err == nil {
+						err = on_data(update_type, b)
+					}
+					update_type = ""
+					if err != nil {
+						return err
+					}
+				} else {
+					data.WriteString(line)
+				}
+			}
+		} else {
+			if line == "KITTY_DATA_START" {
+				started = true
+				update_type = ""
+			}
+		}
+		return nil
+	}
+
+	check_for_line := func() error {
+		if canceled {
+			return nil
+		}
+		s := current_text.String()
+		for {
+			idx := strings.Index(s, "\n")
+			if idx < 0 {
+				break
+			}
+			err = handle_line(s[:idx])
+			if err != nil {
+				return err
+			}
+			s = s[idx+1:]
+		}
+		current_text.Reset()
+		current_text.Grow(4096)
+		if s != "" {
+			current_text.WriteString(s)
+		}
+		return nil
+	}
+
+	lp.OnInitialize = func() (string, error) {
+		pos, chunk_num := 0, 0
+		for {
+			limit := utils.Min(pos+2048, len(data_to_send))
+			if limit <= pos {
+				break
+			}
+			lp.QueueWriteString("\x1bP@kitty-edit|" + strconv.Itoa(chunk_num) + ":")
+			lp.QueueWriteString(data_to_send[pos:limit])
+			lp.QueueWriteString("\x1b\\")
+			chunk_num++
+			pos = limit
+		}
+		lp.QueueWriteString("\x1bP@kitty-edit|\x1b\\")
+		return "", nil
+	}
+
+	lp.OnText = func(text string, from_key_event bool, in_bracketed_paste bool) error {
+		if !from_key_event {
+			current_text.WriteString(text)
+			err = check_for_line()
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+
+	const abort_msg = "\x1bP@kitty-edit|0:abort_signaled=interrupt\x1b\\\x1bP@kitty-edit|\x1b\\"
+
+	lp.OnKeyEvent = func(event *loop.KeyEvent) error {
+		if event.MatchesPressOrRepeat("ctrl+c") || event.MatchesPressOrRepeat("esc") {
+			event.Handled = true
+			canceled = true
+			lp.QueueWriteString(abort_msg)
+			if !started {
+				return tui.Canceled
+			}
+		}
+		return nil
+	}
+
+	err = lp.Run()
+	if err != nil {
+		return
+	}
+	if canceled {
+		return tui.Canceled
+	}
+
+	ds := lp.DeathSignalName()
+	if ds != "" {
+		fmt.Print(abort_msg)
+		if kill_if_signaled {
+			lp.KillIfSignalled()
+			return
+		}
+		return &tui.KilledBySignal{Msg: fmt.Sprint("Killed by signal: ", ds), SignalName: ds}
+	}
+	return
+}
+
+func edit_in_kitty(path string) (err error) {
+	read_file, err := os.Open(path)
+	if err != nil {
+		return fmt.Errorf("Failed to open %s for reading with error: %w", path, err)
+	}
+	defer read_file.Close()
+	var s unix.Stat_t
+	err = unix.Fstat(int(read_file.Fd()), &s)
+	if err != nil {
+		return fmt.Errorf("Failed to stat %s with error: %w", path, err)
+	}
+	if s.Size > 8*1024*1024 {
+		return fmt.Errorf("File size %s is too large for performant editing", humanize.Bytes(uint64(s.Size)))
+	}
+
+	file_data, err := io.ReadAll(read_file)
+	if err != nil {
+		return fmt.Errorf("Failed to read from %s with error: %w", path, err)
+	}
+	read_file.Close()
+	data := strings.Builder{}
+	data.Grow(len(file_data) * 4)
+
+	add := func(key, val string) {
+		if data.Len() > 0 {
+			data.WriteString(",")
+		}
+		data.WriteString(key)
+		data.WriteString("=")
+		data.WriteString(val)
+	}
+	add_encoded := func(key, val string) { add(key, encode(val)) }
+
+	if unix.Access(path, unix.R_OK|unix.W_OK) != nil {
+		return fmt.Errorf("%s is not readable and writeable", path)
+	}
+	cwd, err := os.Getwd()
+	if err != nil {
+		return fmt.Errorf("Failed to get the current working directory with error: %w", err)
+	}
+	add_encoded("cwd", cwd)
+	for _, arg := range os.Args[2:] {
+		add_encoded("a", arg)
+	}
+	add("file_inode", fmt.Sprintf("%d:%d:%d", s.Dev, s.Ino, s.Mtim.Nano()))
+	add_encoded("file_data", utils.UnsafeBytesToString(file_data))
+	fmt.Println("Waiting for editing to be completed, press Esc to abort...")
+	write_data := func(data_type string, rdata []byte) (err error) {
+		err = utils.AtomicWriteFile(path, rdata, fs.FileMode(s.Mode).Perm())
+		if err != nil {
+			err = fmt.Errorf("Failed to write data to %s with error: %w", path, err)
+		}
+		return
+	}
+	err = edit_loop(data.String(), true, write_data)
+	if err != nil {
+		if err == tui.Canceled {
+			return err
+		}
+		return fmt.Errorf("Failed to receive edited file back from terminal with error: %w", err)
+	}
+	return
+}
+
+func EntryPoint(parent *cli.Command) *cli.Command {
+	sc := parent.AddSubCommand(&cli.Command{
+		Name:             "edit-in-kitty",
+		Usage:            "edit-in-kitty [options] file-to-edit",
+		ShortDescription: "Edit a file in a kitty overlay window",
+		HelpText: "Edit the specified file in a kitty overlay window. Works over SSH as well.\n\n" +
+			"For usage instructions see: https://sw.kovidgoyal.net/kitty/shell-integration/#edit-file",
+		Run: func(cmd *cli.Command, args []string) (ret int, err error) {
+			if len(args) == 0 {
+				fmt.Fprintln(os.Stderr, "Usage:", cmd.Usage)
+				return 1, fmt.Errorf("No file to edit specified.")
+			}
+			if len(args) != 1 {
+				fmt.Fprintln(os.Stderr, "Usage:", cmd.Usage)
+				return 1, fmt.Errorf("Only one file to edit must be specified")
+			}
+			err = edit_in_kitty(args[0])
+			return 0, err
+		},
+	})
+	AddCloneSafeOpts(sc)
+	return sc
+}

tools/cmd/tool/main.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"kitty/tools/cli"
 	"kitty/tools/cmd/at"
+	"kitty/tools/cmd/edit_in_kitty"
 	"kitty/tools/cmd/update_self"
 )
 
@@ -18,4 +19,6 @@ func KittyToolEntryPoints(root *cli.Command) {
 	at.EntryPoint(root)
 	// update-self
 	update_self.EntryPoint(root)
+	// edit-in-kitty
+	edit_in_kitty.EntryPoint(root)
 }

tools/cmd/update_self/main.go

@@ -78,7 +78,7 @@ func update_self(version string) (err error) {
 	return unix.Exec(exe, []string{"kitty-tool", "--version"}, os.Environ())
 }
 
-func EntryPoint(root *cli.Command) {
+func EntryPoint(root *cli.Command) *cli.Command {
 	sc := root.AddSubCommand(&cli.Command{
 		Name:             "update-self",
 		Usage:            "update-self [options ...]",
@@ -101,4 +101,5 @@ func EntryPoint(root *cli.Command) {
 		Default: "latest",
 		Help:    "The version to fetch. The special words :code:`latest` and :code:`nightly` fetch the latest stable and nightly release respectively. Other values can be, for example: 0.27.1.",
 	})
+	return sc
 }

tools/utils/atomic-write.go

@@ -0,0 +1,55 @@
+// License: GPLv3 Copyright: 2022, Kovid Goyal, <kovid at kovidgoyal.net>
+
+package utils
+
+import (
+	"fmt"
+	"io/fs"
+	"os"
+	"path/filepath"
+)
+
+var _ = fmt.Print
+
+func AtomicWriteFile(path string, data []byte, perm os.FileMode) (err error) {
+	path, err = filepath.EvalSymlinks(path)
+	if err == nil {
+		path, err = filepath.Abs(path)
+		if err == nil {
+			var f *os.File
+			f, err = os.CreateTemp(filepath.Dir(path), filepath.Base(path))
+			if err == nil {
+				removed := false
+				defer func() {
+					f.Close()
+					if !removed {
+						os.Remove(f.Name())
+					}
+				}()
+				_, err = f.Write(data)
+				if err == nil {
+					err = f.Chmod(perm)
+					if err == nil {
+						err = os.Rename(f.Name(), path)
+						if err == nil {
+							removed = true
+						}
+					}
+				}
+			}
+		}
+	}
+	return
+}
+
+func AtomicUpdateFile(path string, data []byte, perms ...fs.FileMode) (err error) {
+	perm := fs.FileMode(0o666)
+	if len(perms) > 0 {
+		perm = perms[0]
+	}
+	s, err := os.Stat(path)
+	if err == nil {
+		perm = s.Mode().Perm()
+	}
+	return AtomicWriteFile(path, data, perm)
+}