Improved zsh function

README.md

@@ -20,6 +20,39 @@ Install using Go,
 go install github.com/nnutter/git-wt@latest
 ```
 
+## Shell integration
+
+Generate a zsh function that wraps the CLI (default name `wt`):
+
+```bash
+git-wt generate zsh
+# or: git-wt generate zsh --name wt --out $XDG_DATA_HOME/zsh/site-functions --force
+```
+
+Ensure the output directory is on `fpath`, then restart zsh or run `compinit`.
+
+The generated function:
+
+- routes most commands to `git-wt` (`wt create`, `wt list`, `wt prune`, …)
+- provides a shell-only `switch` that `cd`s into a worktree
+- after a successful `wt remove`, `cd`s to the main worktree
+
+```bash
+wt switch main
+wt switch feature/login
+wt create feature/login
+wt remove feature/login   # then cd main
+wt list
+```
+
+If you use [carapace](https://carapace.sh), exclude its built-in `wt` completer (worktrunk) so zsh uses the generated completion instead:
+
+```bash
+export CARAPACE_EXCLUDES=wt
+```
+
+Set this **before** `source <(carapace _carapace)`. You may need `carapace --clear-cache` after changing excludes.
+
 ## Commands
 
 ### `git-wt create <name>`
@@ -77,6 +110,8 @@ When `name` is omitted, removes the managed worktree that contains the current d
 It refuses to remove the main worktree, and refuses dirty or unmerged worktrees by default.
 Use `--force` | `-f` to force (destructive) removal.
 
+When invoked through the shell wrapper (`wt remove`), the shell also switches to the main worktree after a successful removal.
+
 Example:
 
 ```bash
@@ -85,23 +120,25 @@ git-wt remove feature/login
 git-wt remove --force feature/login
 ```
 
-## Typical Flow
+### `git-wt generate zsh`
 
-Checkout [git-cd](https://github.com/nnutter/dotfiles/blob/master/bin/git-cd) to generate shell functions to easily cd to worktrees.
+Generate a zsh wrapper function and completion (see [Shell integration](#shell-integration)).
 
-Create a shell function, `nn`, to switch between the repos under my GitHub path,
+## Typical Flow
 
 ```bash
-git-cd --name nn --repos ~/src/github.com/nnutter
+# once: install wrapper
+git-wt generate zsh
+
+# in a repo
+wt create feature/login
+wt switch feature/login
+# ... work ...
+wt switch main
+wt prune
+# or:
+wt remove feature/login
 ```
 
-Then a typical flow might look like,
-
-```bash
-nn some-repo
-git-wt create feature/login
-...
-nn some-repo.feature.login
-nn some-repo
-git-wt prune
-```
+For jumping between repositories under a path, you can still use something like
+[git-cd](https://github.com/nnutter/dotfiles/blob/master/bin/git-cd).

go.mod

@@ -16,6 +16,7 @@ require (
 	dario.cat/mergo v1.0.2 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/ProtonMail/go-crypto v1.4.1 // indirect
+	github.com/adrg/xdg v0.5.3 // indirect
 	github.com/atotto/clipboard v0.1.4 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/catppuccin/go v0.3.0 // indirect

go.sum

@@ -11,6 +11,8 @@ github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERo
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
 github.com/ProtonMail/go-crypto v1.4.1 h1:9RfcZHqEQUvP8RzecWEUafnZVtEvrBVL9BiF67IQOfM=
 github.com/ProtonMail/go-crypto v1.4.1/go.mod h1:e1OaTyu5SYVrO9gKOEhTc+5UcXtTUa+P3uLudwcgPqo=
+github.com/adrg/xdg v0.5.3 h1:xRnxJXne7+oWDatRhR1JLnvuccuIeCoBu2rtuLqQB78=
+github.com/adrg/xdg v0.5.3/go.mod h1:nlTsY+NNiCBGCK2tpm09vRqfVzrc2fLmXGpBLF0zlTQ=
 github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8=
 github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4=
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio=

internal/gitwt/gitwt.go

@@ -18,7 +18,7 @@ func NewRootCommand() *cobra.Command {
 	rootCommand.AddCommand(NewMigrateCommand())
 	rootCommand.AddCommand(NewPruneCommand())
 	rootCommand.AddCommand(NewRemoveCommand())
-	rootCommand.AddCommand(NewShellCommand())
+	rootCommand.AddCommand(NewGenerateCommand())
 
 	return rootCommand
 }

internal/gitwt/gitwt_generate.go

@@ -0,0 +1,27 @@
+package gitwt
+
+import (
+	"github.com/spf13/cobra"
+)
+
+type generateCommandOptions struct{}
+
+func NewGenerateCommand() *cobra.Command {
+	options := &generateCommandOptions{}
+
+	command := &cobra.Command{
+		Use:   `generate`,
+		Short: `Generate shell integration wrapping git-wt`,
+		Args:  cobra.NoArgs,
+		RunE:  options.Execute,
+	}
+	command.CompletionOptions.HiddenDefaultCmd = true
+
+	command.AddCommand(NewZshCommand())
+
+	return command
+}
+
+func (x *generateCommandOptions) Execute(command *cobra.Command, args []string) error {
+	return command.Help()
+}

internal/gitwt/gitwt_generate_zsh.go

@@ -0,0 +1,194 @@
+package gitwt
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/spf13/cobra"
+)
+
+type zshCommandOptions struct {
+	name  string
+	out   string
+	force bool
+}
+
+func NewZshCommand() *cobra.Command {
+	options := &zshCommandOptions{}
+
+	command := &cobra.Command{
+		Use:   `zsh`,
+		Short: `Generate zsh function and completion wrapping git-wt`,
+		Args:  cobra.NoArgs,
+		RunE:  options.Execute,
+	}
+
+	command.Flags().StringVarP(&options.name, `name`, `n`, `wt`, `name of generated zsh function`)
+	command.Flags().StringVarP(&options.out, `out`, `o`, xdgDataHome()+`/zsh/site-functions`, `output directory for generated files`)
+	command.Flags().BoolVarP(&options.force, `force`, `f`, false, `overwrite existing generated files`)
+
+	return command
+}
+
+func xdgDataHome() string {
+	if xdg := os.Getenv(`XDG_DATA_HOME`); xdg != `` {
+		return xdg
+	}
+	return filepath.Join(os.Getenv(`HOME`), `.local`, `share`)
+}
+
+func (x *zshCommandOptions) Execute(command *cobra.Command, args []string) error {
+	outDir := x.out
+
+	if err := os.MkdirAll(outDir, 0o755); err != nil {
+		return fmt.Errorf(`create output directory: %w`, err)
+	}
+
+	functionPath := filepath.Join(outDir, x.name)
+	completionPath := filepath.Join(outDir, `_`+x.name)
+
+	if !x.force {
+		if _, err := os.Stat(functionPath); err == nil {
+			return fmt.Errorf(`function file %q already exists; use --force to overwrite`, functionPath)
+		}
+		if _, err := os.Stat(completionPath); err == nil {
+			return fmt.Errorf(`completion file %q already exists; use --force to overwrite`, completionPath)
+		}
+	}
+
+	if err := x.writeFunctionFile(functionPath); err != nil {
+		return err
+	}
+
+	if err := x.writeCompletionFile(completionPath); err != nil {
+		return err
+	}
+
+	fmt.Fprintf(command.ErrOrStderr(), "Wrote %s\n", functionPath)
+	fmt.Fprintf(command.ErrOrStderr(), "Wrote %s\n", completionPath)
+	fmt.Fprintf(command.ErrOrStderr(), "Ensure %s is on fpath. Restart zsh or run compinit if completion does not appear immediately.\n", outDir)
+
+	return nil
+}
+
+func (x *zshCommandOptions) writeFunctionFile(target string) error {
+	content := `# Generated by git-wt generate zsh
+# DO NOT EDIT. Regenerate with git-wt generate zsh
+
+` + x.name + `() {
+    case "$1" in
+    switch)
+        shift
+        if [[ -z "$1" ]]; then
+            echo "Usage: ` + x.name + ` switch <worktree>" >&2
+            return 1
+        fi
+
+        local main_dir
+        main_dir=$(git worktree list --porcelain | head -n1 | sed "s/^worktree //")
+        if [[ -z "$main_dir" ]]; then
+            echo "Main worktree not found" >&2
+            return 1
+        fi
+
+        local parent_dir=$(dirname "$main_dir")
+        local repo_name=$(basename "$main_dir")
+        local arg=$1
+
+        if [[ $arg == $repo_name.* ]]; then
+            arg=${arg#$repo_name.}
+        fi
+
+        case "$arg" in
+        main)
+            if [[ $(pwd) == "$main_dir" ]]; then
+                echo "Already in main worktree"
+                return 0
+            fi
+            if ! [[ -d "$main_dir" ]]; then
+                echo "Main worktree not found" >&2
+                return 1
+            fi
+            cd "$main_dir"
+            ;;
+        *)
+            local target_dir=$parent_dir/${repo_name}.${arg//\//.}
+            if [[ $(pwd) == "$target_dir" ]]; then
+                echo "Already in $arg"
+                return 0
+            fi
+            if ! [[ -d "$target_dir" ]]; then
+                echo "Worktree $arg not found at $target_dir" >&2
+                return 1
+            fi
+            cd "$target_dir"
+            ;;
+        esac
+        ;;
+    remove)
+        local main_dir
+        main_dir=$(git worktree list --porcelain | head -n1 | sed "s/^worktree //")
+        if [[ -z "$main_dir" ]]; then
+            echo "Main worktree not found" >&2
+            return 1
+        fi
+        command git-wt "$@" || return $?
+        cd "$main_dir"
+        ;;
+    *)
+        command git-wt "$@"
+        ;;
+    esac
+}
+`
+
+	return os.WriteFile(target, []byte(content), 0o644)
+}
+
+func (x *zshCommandOptions) writeCompletionFile(target string) error {
+	content := `#compdef ` + x.name + `
+
+# Generated by git-wt generate zsh
+# DO NOT EDIT. Regenerate with git-wt generate zsh
+
+_` + x.name + `() {
+    local -a subcommands
+    subcommands=(
+        'create:Create a managed Git worktree'
+        'list:List managed Git worktrees'
+        'migrate:Bring existing worktrees under management'
+        'prune:Remove clean merged managed worktrees'
+        'remove:Remove a managed Git worktree'
+        'generate:Generate shell integration'
+        'switch:Switch to a worktree'
+    )
+
+    if (( CURRENT == 2 )); then
+        _describe 'command' subcommands
+        return
+    fi
+
+    case $words[2] in
+    switch|remove)
+        if ! git rev-parse --is-inside-work-tree 1>/dev/null 2>/dev/null; then
+            return 1
+        fi
+
+        local main_dir
+        main_dir=$(git worktree list --porcelain | head -n1 | sed "s/^worktree //")
+
+        local -a worktrees
+        if [[ $words[2] == switch ]]; then
+            worktrees=(main)
+        fi
+        worktrees+=($(git worktree list --porcelain 2>/dev/null | grep '^worktree ' | tail -n +2 | sed 's|^worktree '"$main_dir"'\.||'))
+
+        _describe 'worktrees' worktrees
+        ;;
+    esac
+}
+`
+
+	return os.WriteFile(target, []byte(content), 0o644)
+}