From 55925064567a852d0a96c6e32833836421849c81 Mon Sep 17 00:00:00 2001
From: Ajeet D'Souza <98ajeet@gmail.com>
Date: Wed, 21 Feb 2024 01:45:20 +0530
Subject: [PATCH] Improve zsh completions (#725)

---
 CHANGELOG.md      |  6 ++++
 templates/zsh.txt | 85 +++++++++++++++++++++++++++--------------------
 2 files changed, 55 insertions(+), 36 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 65f3392..2db5f58 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+### Changed
+
+- zsh: improved Space-Tab completions.
+
 ## [0.9.3] - 2024-02-13
 
 ### Added
diff --git a/templates/zsh.txt b/templates/zsh.txt
index 20ddd5c..a45c710 100644
--- a/templates/zsh.txt
+++ b/templates/zsh.txt
@@ -52,8 +52,6 @@ fi
 # When using zoxide with --no-cmd, alias these internal functions as desired.
 #
 
-__zoxide_z_prefix='z#'
-
 # Jump to a directory using only keywords.
 function __zoxide_z() {
     # shellcheck disable=SC2199
@@ -61,15 +59,10 @@ function __zoxide_z() {
         __zoxide_cd ~
     elif [[ "$#" -eq 1 ]] && { [[ -d "$1" ]] || [[ "$1" = '-' ]] || [[ "$1" =~ ^[-+][0-9]$ ]]; }; then
         __zoxide_cd "$1"
-    elif [[ "$@[-1]" == "${__zoxide_z_prefix}"?* ]]; then
-        # shellcheck disable=SC2124
-        \builtin local result="${@[-1]}"
-        __zoxide_cd "{{ "${result:${#__zoxide_z_prefix}}" }}"
     else
         \builtin local result
         # shellcheck disable=SC2312
-        result="$(\command zoxide query --exclude "$(__zoxide_pwd)" -- "$@")" &&
-            __zoxide_cd "${result}"
+        result="$(\command zoxide query --exclude "$(__zoxide_pwd)" -- "$@")" && __zoxide_cd "${result}"
     fi
 }
 
@@ -79,32 +72,6 @@ function __zoxide_zi() {
     result="$(\command zoxide query --interactive -- "$@")" && __zoxide_cd "${result}"
 }
 
-# Completions.
-if [[ -o zle ]]; then
-    function __zoxide_z_complete() {
-        # Only show completions when the cursor is at the end of the line.
-        # shellcheck disable=SC2154
-        [[ "{{ "${#words[@]}" }}" -eq "${CURRENT}" ]] || return 0
-
-        if [[ "{{ "${#words[@]}" }}" -eq 2 ]]; then
-            _files -/
-        elif [[ "${words[-1]}" == '' ]] && [[ "${words[-2]}" != "${__zoxide_z_prefix}"?* ]]; then
-            \builtin local result
-            # shellcheck disable=SC2086,SC2312
-            if result="$(\command zoxide query --exclude "$(__zoxide_pwd)" --interactive -- ${words[2,-1]})"; then
-                result="${__zoxide_z_prefix}${result}"
-                # shellcheck disable=SC2296
-                compadd -Q "${(q-)result}"
-            fi
-            \builtin printf '\e[5n'
-        fi
-        return 0
-    }
-
-    \builtin bindkey '\e[0n' 'reset-prompt'
-    [[ "${+functions[compdef]}" -ne 0 ]] && \compdef __zoxide_z_complete __zoxide_z
-fi
-
 {{ section }}
 # Commands for zoxide. Disable these using --no-cmd.
 #
@@ -112,8 +79,54 @@ fi
 {%- match cmd %}
 {%- when Some with (cmd) %}
 
-\builtin alias {{cmd}}=__zoxide_z
-\builtin alias {{cmd}}i=__zoxide_zi
+function {{ cmd }}() {
+    __zoxide_z "$@"
+}
+
+function {{ cmd }}i() {
+    __zoxide_zi "$@"
+}
+
+# Completions.
+if [[ -o zle ]]; then
+    __zoxide_result=''
+
+    function __zoxide_z_complete() {
+        # Only show completions when the cursor is at the end of the line.
+        # shellcheck disable=SC2154
+        [[ "{{ "${#words[@]}" }}" -eq "${CURRENT}" ]] || return 0
+
+        if [[ "{{ "${#words[@]}" }}" -eq 2 ]]; then
+            # Show completions for local directories.
+            _files -/
+        elif [[ "${words[-1]}" == '' ]]; then
+            # Show completions for Space-Tab.
+            # shellcheck disable=SC2086
+            __zoxide_result="$(\command zoxide query --exclude "$(__zoxide_pwd || \builtin true)" --interactive -- ${words[2,-1]})" || __zoxide_result=''
+            # Sends '\e[0n' to console input.
+            \builtin printf '\e[5n'
+        fi
+
+        # Report that the completion was successful, so that we don't fall back
+        # to another completion function.
+        return 0
+    }
+
+    function __zoxide_z_complete_helper() {
+        if [[ -n "${__zoxide_result}" ]]; then
+            # shellcheck disable=SC2034,SC2296
+            BUFFER="{{ cmd }} ${(q-)__zoxide_result}"
+            \builtin zle reset-prompt
+            \builtin zle accept-line
+        else
+            \builtin zle reset-prompt
+        fi
+    }
+    \builtin zle -N __zoxide_z_complete_helper
+
+    \builtin bindkey '\e[0n' '__zoxide_z_complete_helper'
+    [[ "${+functions[compdef]}" -ne 0 ]] && \compdef __zoxide_z_complete {{ cmd }}
+fi
 
 {%- when None %}