Add worktree {add,remove} <branch> command

Wrap `git worktree {add,remove} <path> <branch>` to streamline working
with multiple worktrees.

worktree/_worktree

@@ -0,0 +1,50 @@
+#compdef worktree
+
+__worktree_branches() {
+  local -a branches
+  branches=(
+    ${(fo)"$(git branch --format='%(refname:short)' 2>/dev/null)"}
+    ${${${${(fo)"$(git for-each-ref --format='%(refname)' refs/remotes 2>/dev/null)"}#refs/remotes/}#*/}:#HEAD}
+  )
+  _describe 'branch' branches
+}
+
+__worktree_active_branches() {
+  local -a branches
+  branches=(${(fo)"$(git worktree list --porcelain 2>/dev/null | \
+    awk '/^worktree /{ wt++ } wt>1 && /^branch refs\/heads\//{ sub(/^branch refs\/heads\//, ""); print }')"})
+  _describe 'branch' branches
+}
+
+_worktree() {
+  local context curcontext="$curcontext" state line
+  typeset -A opt_args
+
+  _arguments -C \
+    '1: :->cmd' \
+    '*:: :->args'
+
+  case $state in
+    (cmd)
+      local commands; commands=(
+        'add:Create a worktree for a branch'
+        'remove:Remove a worktree by branch'
+        'rm:Remove a worktree by branch'
+      )
+      _describe -t commands 'worktree command' commands "$@"
+      ;;
+    (args)
+      curcontext="${curcontext%:*:*}:worktree-cmd-$words[1]:"
+      case $line[1] in
+        (add)
+          _arguments '1:: :__worktree_branches'
+          ;;
+        (rm|remove)
+          _arguments '1:: :__worktree_active_branches'
+          ;;
+      esac
+      ;;
+  esac
+}
+
+_worktree "$@"

worktree/worktree.plugin.zsh

@@ -0,0 +1,66 @@
+worktree() {
+  if [[ "$1" == "" ]]; then
+    echo "usage: worktree {add,remove} <branch>"
+    return 1
+  elif [[ "$1" == "-h" ]] || [[ "$1" == "--help" ]]; then
+    echo "usage: worktree {add,remove} <branch>
+
+Manage git worktrees by branch name."
+    return 0
+  fi
+
+  local cmd=$1
+  local branch=$2
+
+  if [[ -z "$branch" ]]; then
+    print -P "%F{red}error:%f missing <branch> argument"
+    return 1
+  fi
+
+  if ! git rev-parse --git-dir &>/dev/null; then
+    print -P "%F{red}error:%f not a git repository"
+    return 1
+  fi
+
+  if [[ "$(git rev-parse --is-inside-work-tree 2>/dev/null)" == "true" ]] && \
+     [[ "$(git rev-parse --absolute-git-dir)" != "$(git rev-parse --show-toplevel)/.git" ]]
+  then
+    print -P "%F{red}error:%f must be run from the main clone, not a worktree"
+    return 1
+  fi
+
+  case $cmd in
+    add)
+      local root=$(git rev-parse --show-toplevel)
+      local wt_dest=${root:h}/${root:t}@${branch}
+      git worktree add "$wt_dest" "$branch"
+      ;;
+    rm|remove)
+      local wt_path
+      wt_path=$(git worktree list --porcelain | awk -v b="$branch" \
+        '/^worktree /{ sub(/^worktree /, ""); p=$0 } /^branch refs\/heads\//{ sub(/^branch refs\/heads\//, ""); if($0==b) print p }')
+      if [[ -z "$wt_path" ]]; then
+        print -P "%F{red}error:%f no worktree for branch: $branch"
+        return 1
+      fi
+      if [[ -f "$wt_path/BUILD" ]] || [[ -f "$wt_path/WORKSPACE" ]]; then
+        if (( $+commands[bazel] )); then
+          bazel --output_base="$wt_path" clean --expunge
+        elif [[ -x "$wt_path/bazelw" ]]; then
+          "$wt_path/bazelw" clean --expunge
+        fi
+      fi
+      git worktree remove "$wt_path" || return 1
+      local root=$(git rev-parse --show-toplevel)
+      local parent=${wt_path:h}
+      while [[ "$parent" != "${root:h}" ]] && [[ -d "$parent" ]]; do
+        rmdir "$parent" 2>/dev/null || break
+        parent=${parent:h}
+      done
+      ;;
+    *)
+      print -P "%F{red}error:%f unknown command: $cmd"
+      return 1
+      ;;
+  esac
+}

zshrc

@@ -50,6 +50,9 @@ source-plugin session
 # Note taking commands
 source-plugin notes
 
+# git worktree helper
+source-plugin worktree
+
 # Remove duplicates from history
 setopt histignoredups