Suggesting more zsh auto-suggestions

The mechanism behind tab

In zsh, there is the concept of a widget, which are small, often contained actions performed by the zsh line editor (ZLE). The ZLE is responsible for letting you issue commands in the shell in an interactive way; it's the line editor that lets type into the shell. Simple enough.

Widgets are just regular zsh functions. You can write and register one easily enough:

#!/bin/zsh
function _ouch() {
    echo -n "Ouch! Don't do that!"
}

# use -N to register a new widget
zle -N ouch-widget _ouch

We can invoke this widget using the execute-named-cmd widget. By default (emacs mode), this is bound to alt-x:

You can bind arbitrary widgets to certain keystrokes (escape codes), and have them perform some action. There's a list of builtin widgets that you can browse, for example: end-of-line will for most be be bound to ctrl-e, or up-history is likely the up arrow key. These are all defined as little scripts that the ZLE executes. That means they have access to variables that only exist within the scope of the ZLE, such as $BUFFER, which is the current line being edited (what you are typing into the prompt), or $CURSOR which is the position of the caret in the buffer. There's a full list available in the docs. You can likely already piece together how tab completion works as a widget: find the current word being typed with ${BUFFER[$CURSOR]}, and then do a look up in a dictionary of possible completions. A good dictionary would be context specific, for a buffer that starts with git will have a different dictionary from a command that starts with rm. You can generalize this to subcommands or other information present (files in the CWD, known hosts from the hosts file, and so on) .

To make this more manageable, zsh has a full completion system that tries to make this smooth, fast, and highly customizable. If you use zsh, you likely have the following in your .zshrc:

autoload -U compinit
compinit

which loads and initializes the completion system with all of its associated widgets, and rebinds tab from the expand-or-complete widget to complete-word. This system can be configured with zstyle in a myriad of different ways; it can present the completion to you as a simple list under the current line, or, say, as an interactive menu, with extra detailed context and completion groups. That's all of that is a little beyond what I want to focus on in this post however. It's sufficient to know that compinit will enable all the things that builds those lookup dictionaries in a manner that can be telescopically specific, and it does all this when you hit tab.

The problem that the suggestions system has is that there's no other way to query the completion dictionaries other than by pressing tab.

So, if you want to modify the completion system or somehow reap the dictionary it sows, there's no function that you can call and have it return an array of suggestions. This is because the underlying widgets write directly into buffers and $POSTDISPLAY variables that get written to the shell's display and don't touch file descriptors. The data structure that keeps the dictionary remains elusive and out of reach.

Injecting completion behaviour to reap completion

But if the completion dictionaries were unreachable, how do tools like zsh-autosuggestions know anything about what to suggest? Well, if we look at what their code is doing we quickly discover the hack: spawn a new shell, give it the same line buffer, then send it the tab key press event, and parse what it outputs into a suggestion dictionary.

We don't need a whole new shell to do that though, we just need another ZLE. Fortunately, zsh has a module called zsh/zpty that lets us handle pseudo-terminals, designed to enable a degree of asynchronous execution in zsh: you can use zpty to wrap interactive commands and script the interaction.

Let's write an example. We'll spawn a zpty, send the tab keypress, and read the resulting output buffer to get the list of suggestions. This is slightly less complicated but a little hackier than obtaining only the top suggestion, but it serves to illustrate the approach.

#!/bin/zsh
MY_PTY="some-unique-name-to-avoid-clobbering"

# load zpty and compinit
zmodload zsh/zpty 2>/dev/null || return
autoload -Uz compinit ; compinit

get_completion() {
    # setup the zpty with a buffer
    zpty $MY_PTY capture_completion "\$1"
    # write to the PTY
    zpty -w $MY_PTY $'\t'
    {
        local out last=() suggestions=()
        # read first line; this will be the buffer
        zpty -r $MY_PTY out
        while zpty -r $MY_PTY out; do
            # append what we gathered in the last loop. this allows us to avoid
            # reading the final line into the suggestion buffer, since this is
            # just a repeat of the input
            suggestions+=$last
            # strip off carriage returns and new lines
            out="${out//[$'\n\r']}"
            last=(${(z)out})
        done
        # remove nulls and empties
        suggestions=(${(@)suggestions:#})
        echo -n "$suggestions"
    } always {
        # destroy the pty
        zpty -d $MY_PTY
    }
}

That's the skeleton of what we're going to be doing. Most of this is just reaping the output. We create a new zpty that runs a (yet to be defined) capture_completion function and forwards the argument from get_completion (which will be the string we want to get completions for). To read the output, we skip the first line, as this is the $BUFFER. Everything else is then the complete completion suggestions; these are words split by some number of spaces, so we can use a zsh parameter expansion flag to parse this into an array, trimming all whitespaces. The very last line of the output buffer is the input, which is not a completion, so we remove it. Because the input can be multiple (space separated) words, we can't just remove the last element from $suggestions, and instead use a $last array to stagger appending the output.

We'll invoke this function like so

get_completion "git st"

and it will output all the things that could complete this line as generated by the git shell completion scripts.

The next stage is to implement the thing that will run inside the zpty, i.e. the capture_completion function. This function will setup our stage as we ask our actors to enter: we'll use zstyle to mask out all the additional information zsh likes to give us when we ask for completion (description of the flags, for example), and to setup a hook into the completion system. This hook is needed to stop the completion system inserting the suggestion, and force it to display the options to the PTY instead. For example, if we had git statu then this would unambiguously complete to git status, and hitting tab will just complete the word. This needs to be disabled, so that the suggestion is instead printed to the screen. The hook is inserted by removing vared in the compstate before handing off to the completion function.

capture_completion() {
    # disable all description text in completion lists so we just get the words
    zstyle ':completion:*' format ''
    zstyle ':completion:*:descriptions' format ''
    zstyle ':completion:*' group-name ''
    zstyle ':completion:*:messages' format ''
    zstyle ':completion:*:warnings' format ''
    zstyle ':completion:*' verbose no

    # hook in a custom completion widget
    bindkey '^I' capture-completion

    # for $functions, so that we can change the `_complete` function
    zmodload zsh/parameter 2>/dev/null || return
    autoload +X _complete

    functions[_original_complete]=$functions[_complete]
    function _complete() {
        # force: do not edit the current command line
        unset 'compstate[vared]'
        # call the original tab completion function
        _original_complete "$@"
    }

    # put contents of $1 in the buffer so we're ready to accept the tab and
    # capture completions
    vared 1
}

The last thing to do is to define the capture-completion widget. This will just call the complete-word widget so that all the suggestions are printed to the PTY.

_capture_completion_widget() {
    # call the original tab completion function
    zle complete-word
}

# register our widget
zle -N capture-completion _capture_completion_widget

That's the full round trip for generating a list of completion suggestions:

$ get_completion "git s"
sd sw shortlog sparse-checkout status subtree show stash submodule switch
show-branch send-email svn symbolic-ref show-index show-ref send-pack shell
stripspace

ZSH auto suggestions

The zsh-autosuggestions plugin works with more or less the above mechanism. Indeed, the above example is a modified version of the zsh-autosuggestions source. It differs slightly in the current version in that it only captures one word, and does so by disabling the list feature in compstate, and forces compstate to always write the first completion into the buffer. The suggestion it then fetches includes the buffer. For example, git st will be completed not just as stash or ash, but becomes the full line git stash. I believe this is to handle edge cases that my above example would fail on when there are spaces in arguments, but I am not entirely sure.

Once the suggestions have been gathered, we then need to display them. Adding them to the $BUFFER would certainly work, but now we have fundamentally changed the command line and not just displayed some new information. There is instead the $POSTDISPLAY variable which we can assign to, which will draw text after the buffer without changing the content of the buffer. Additionally, text in $POSTDISPLAY cannot be edited directly by the user.

Let's implement a naive version of this with the list of suggestions we've been able to gather. We'll define a new widget which will forward the buffer to get_completion, and then write to the POSTDISPLAY to show the results.

function _list_suggestions() {
    # return variable
    local completions=$(get_completion "$BUFFER")
    local -a suggestions=()
    for word in $(echo $completions \
        | sed -r "s/\x1B\[([0-9]{1,3}(;[0-9]{1,2};?)?)?[mGKJ]//g"); do
        suggestions+=("$word")
    done
    # put curly brackets around so it's clear what we're adding
    POSTDISPLAY="{${(j| |)suggestions}}"
}

zle -N list-suggestions _list_suggestions

# bind alt-k to activate our widget
bindkey "\ek" list-suggestions

There are many ways to skin a cat, but we'll use sed to strip off any control sequences that are still present in the completions list (I am unsure where they come from but there they are). We then put these pristine suggestions in to the $suggestions array, and join each element with a space using another zsh parameter expansion flag.

Now we get something like this when we hit alt-k:

There are two problems however: visual feedback and latency. We'll address the prior, but the latter requires a more complex asynchronous solution within zsh. In essence, some commands run a fairly heavy job to generate their completions, and we're blocking until that has finished. The zsh-autosuggestions plugin circumvents this by running the completion widget asynchronously using a seperate process, reaping the output from a file descriptor, and therefore only updates $POSTDISPLAY when the completion is finished without blocking. It's complex but works well, and I will likely write a future post with the fundamentals, but will avoid discussing this today. We'll just have to put up with the latency in this example for now.

The visual feedback we can easily solve using the region_highlight parameter. This work by appending the start and end indices we want to highlight with the attributes we want to apply to that region.

function _show_suggestion() {
    # unset the region highlights from the last completion
    unset region_highlight
    zle list-suggestions
    # apply to just the $POSTDISPLAY region, so get the length of the buffer
    # and add it it the length of the post diplay
    local start=$#BUFFER
    local end=$(($start + $#POSTDISPLAY))
    region_highlight="$start $end fg=yellow"
}

zle -N show-suggestions _show_suggestion

# re-bind alt-k to activate our show widget
bindkey "\ek" show-suggestions

The last thing to do to make this fully functional is to trigger the widget on each key press. The zsh-autosuggestions plugin has a long list of widgets that it overrides to bind its own widgets to. We'll just be very simple with our toy example and bind to the character insertion widget:

function suggestions_start() {
    local widgets_to_bind=(
        self-insert
    )

    local widget
    for widget ($widgets_to_bind); do
        eval "_suggestion_${(q)widget}() {
            zle show-suggestions
            zle .${(q)widget}
        }"
        # Create the bound widget
        zle -N "$widget" "_suggestion_${(q)widget}"
    done
}

suggestions_start

Great! Now all we need to do is fix the latency problems. But let's not reinvent the wheel. Now that we have a firm understanding of how this all works, we'll instead just modify the source of zsh-autosuggestions with what we've learned.

Auto-suggesting more than one suggestion

I've made various modifications to fork of zsh-autosuggestions to implement this behaviour. It doesn't work well with the existing features (e.g. breaks history mode), so I doubt I will try to merge it upstream anytime soon, but for my own purposes it works well.

The key changes I made were as follows:

• Store all of the suggestions in a $GLOBAL_MATCHES variable, so that they can be accessed at any point.
• When the user types, zsh-autosuggestions calls a modify hook to check if the currently typed character matches with the current suggestion or not. If it does, it will skip the call to fetch the suggestions, and just use what it already has. It looked as if zsh-autosuggestions was "storing state" in the $POSTDISPLAY variable, so when it would do this check, it just does a string compare against that variable. With multiple suggestions, however, there is now the need to perform that check for each suggestion. That's a feature I added, so it will look at the current $GLOBAL_MATCHES and will avoid the round trip to find new suggestions if the user is typing any one of them.
• Avoid filling the terminal: some commands have 100s of suggestions, so to avoid blowing the terminal up when a user starts typing that kind of command, I limit the number of characters it will suggest to some hard-coded value. I chose 150 characters arbitrarily, and made it insert a ... to indicate there are more suggestions that it's not showing.
• Remove zsh: do you want to see all X possibilities. Basically if zsh thinks your terminal is too small it will, instead of printing the suggestions, show a helpful little message instead that asks the user if they want to see all of the options or not. This avoid flooding the terminal. But I've already addressed that with the character limit, so we want to get rid of it. The easiest thing to do is to pretend the terminal is enormous when we ask zsh for suggestions. We can do that in the PTY using COLUMNS=1000 ; LINES=1000, and now zsh will always display all of the options for us to collect.

A bed full of bugs

My code is messy. I am using it on all of my machines despite that, and encountered many, many bugs, most of which I think I've fixed. I haven't found too many places where it doesn't work, other than when certain control characters are present in the buffer, but I am struggling to reproduce that consistently enough to actually debug it.

The one major bug that I do know about but haven't had the time to fix yet is that the suggestions can get a little muddled if you start descending file paths with / in them.

If you want to use it too, I leave a link to the code below. Use at your own risk, and so on. If you find more bugs, please do let me know, but I make no promises to fix them <3

• fjebaker/more-zsh-autosuggestions

Notes

1. There's probably a much easier way to do all the zpty output parsing. In zsh, there is the list-choices widget that the completion system provides, which does pretty much exactly what it says on the tin: it lists the possible completions to the PTY. This could be used instead of having to unset vared in the compstate, and parse the buffer in addition to the suggestions.