Automatically setup shell integration for ZSH

kitty/child.py

@@ -227,6 +227,11 @@ class Child:
             tdir = checked_terminfo_dir()
             if tdir:
                 env['TERMINFO'] = tdir
+            opts = fast_data_types.get_options()
+            if opts.shell_integration != 'disabled':
+                from .shell_integration import get_supported_shell_name
+                if get_supported_shell_name(self.argv[0]):
+                    env['KITTY_SHELL_INTEGRATION'] = opts.shell_integration
         return env
 
     def fork(self) -> Optional[int]:

kitty/main.py

@@ -151,6 +151,9 @@ def _run_app(opts: Options, args: CLIOptions, bad_lines: Sequence[BadLine] = ())
     if not is_wayland() and not is_macos:  # no window icons on wayland
         set_x11_window_icon()
     load_shader_programs.use_selection_fg = opts.selection_foreground is not None
+    if opts.shell_integration != 'disabled':
+        from .shell_integration import setup_shell_integration
+        setup_shell_integration()
     with cached_values_for(run_app.cached_values_name) as cached_values:
         with startup_notification_handler(extra_callback=run_app.first_window_callback) as pre_show_callback:
             window_id = create_os_window(

kitty/shell_integration.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+# vim:fileencoding=utf-8
+# License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>
+
+
+import os
+from typing import Optional
+
+from .constants import shell_integration_dir
+from .fast_data_types import get_options
+from .types import run_once
+from .utils import log_error, resolved_shell
+
+SUPPORTED_SHELLS = ('zsh',)
+posix_template = '''
+# BEGIN_KITTY_SHELL_INTEGRATION
+[[ -a {path} ]] && source {path}
+# END_KITTY_SHELL_INTEGRATION
+'''
+
+
+def setup_integration(shell_name: str, rc_path: str, template: str = posix_template) -> None:
+    import re
+    rc_path = os.path.realpath(rc_path)
+    if not os.access(rc_path, os.W_OK, effective_ids=os.access in os.supports_effective_ids):
+        return
+    with open(rc_path) as f:
+        rc = f.read()
+    home = os.path.expanduser('~') + '/'
+    path = os.path.join(shell_integration_dir, f'kitty.{shell_name}')
+    if path.startswith(home):
+        path = '$HOME/' + path[len(home):]
+    integration = template.format(path=f'"{path}"')
+    newrc = re.sub(
+        r'^# BEGIN_KITTY_SHELL_INTEGRATION.+?^# END_KITTY_SHELL_INTEGRATION',
+        '', rc, flags=re.DOTALL | re.MULTILINE)
+    newrc = newrc.rstrip() + '\n\n' + integration
+    if newrc != rc:
+        tmp = rc_path + '_ksi_tmp'
+        with open(tmp, 'w') as f:
+            f.write(newrc)
+        os.rename(tmp, rc_path)
+
+
+def setup_zsh_integration() -> None:
+    base = os.environ.get('ZDOTDIR', os.path.expanduser('~'))
+    rc = os.path.join(base, '.zshrc')
+    setup_integration('zsh', rc)
+
+
+def get_supported_shell_name(path: str) -> Optional[str]:
+    name = os.path.basename(path).split('.')[0].lower()
+    if name in SUPPORTED_SHELLS:
+        return name
+
+
+@run_once
+def setup_shell_integration() -> None:
+    opts = get_options()
+    q = opts.shell_integration.split()
+    if opts.shell_integration == 'disabled' or 'no-rc' in q:
+        return
+    shell = get_supported_shell_name(resolved_shell(opts)[0])
+    if shell is None:
+        return
+    func = {'zsh': setup_zsh_integration}[shell]
+    try:
+        func()
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        log_error(f'Failed to setup shell integration for: {shell}')

shell-integration/kitty.zsh

@@ -1,14 +1,14 @@
 () {
     if [[ ! -o interactive ]]; then return; fi
-    if [[ -z "$kitty_shell_integration" ]]; then return; fi
+    if [[ -z "$KITTY_SHELL_INTEGRATION" ]]; then return; fi
     typeset -g -A _ksi_prompt=([state]='first-run' [cursor]='y' [title]='y' [mark]='y' [complete]='y')
-    for i in ${=kitty_shell_integration}; do
+    for i in ${=KITTY_SHELL_INTEGRATION}; do
         if [[ "$i" == "no-cursor" ]]; then _ksi_prompt[cursor]='n'; fi
         if [[ "$i" == "no-title" ]]; then _ksi_prompt[title]='n'; fi
         if [[ "$i" == "no-prompt-mark" ]]; then _ksi_prompt[mark]='n'; fi
         if [[ "$i" == "no-complete" ]]; then _ksi_prompt[complete]='n'; fi
     done
-    unset kitty_shell_integration
+    unset KITTY_SHELL_INTEGRATION
 
     function _ksi_debug_print() {
         # print a line to STDOUT of parent kitty process