Add worktree cd subcommand & --force flag to remove

worktree/_worktree

@@ -16,6 +16,13 @@ __worktree_active_branches() {
   _describe 'branch' branches
 }
 
+__worktree_all_branches() {
+  local -a branches
+  branches=(${(fo)"$(git worktree list --porcelain 2>/dev/null | \
+    awk '/^branch refs\/heads\//{ sub(/^branch refs\/heads\//, ""); print }')"})
+  _describe 'branch' branches
+}
+
 _worktree() {
   local context curcontext="$curcontext" state line
   typeset -A opt_args
@@ -28,6 +35,7 @@ _worktree() {
     (cmd)
       local commands; commands=(
         'add:Create a worktree for a branch'
+        'cd:Change directory to a worktree by branch'
         'list:List managed worktrees'
         'ls:List managed worktrees'
         'remove:Remove a worktree by branch'
@@ -42,8 +50,13 @@ _worktree() {
         (add)
           _arguments '1:: :__worktree_branches'
           ;;
+        (cd)
+          _arguments '1:: :__worktree_all_branches'
+          ;;
         (rm|remove)
-          _arguments '*:: :__worktree_active_branches'
+          _arguments \
+            '(-f --force)'{-f,--force}'[Force removal even with modified or untracked files]' \
+            '*:: :__worktree_active_branches'
           ;;
         (sync)
           _arguments '*:: :__worktree_active_branches'

worktree/worktree.plugin.zsh

@@ -1,9 +1,9 @@
 worktree() {
   if [[ "$1" == "" ]]; then
-    echo "usage: worktree {add,list,remove,sync} <branch> [<branch>...]"
+    echo "usage: worktree {add,cd,list,remove,sync} <branch> [<branch>...]"
     return 1
   elif [[ "$1" == "-h" ]] || [[ "$1" == "--help" ]]; then
-    echo "usage: worktree {add,list,remove,sync} <branch> [<branch>...]
+    echo "usage: worktree {add,cd,list,remove,sync} <branch> [<branch>...]
 
 Manage git worktrees by branch name."
     return 0
@@ -17,7 +17,8 @@ Manage git worktrees by branch name."
     return 1
   fi
 
-  if [[ "$(git rev-parse --is-inside-work-tree 2>/dev/null)" == "true" ]] && \
+  if [[ "$cmd" != "cd" ]] && \
+     [[ "$(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"
@@ -44,6 +45,20 @@ Manage git worktrees by branch name."
         _autoenv_authorized $wt_dest/.exit yes
       fi
       ;;
+    cd)
+      if [[ -z "$branch" ]]; then
+        print -P "%F{red}error:%f missing <branch> argument"
+        return 1
+      fi
+      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
+      cd "$wt_path"
+      ;;
     list|ls)
       local root=$(git rev-parse --show-toplevel)
       local prefix="${root:h}/${root:t}@"
@@ -53,11 +68,19 @@ Manage git worktrees by branch name."
       '
       ;;
     rm|remove)
-      if [[ -z "$branch" ]]; then
+      local force=0
+      local branches=()
+      local arg
+      for arg in "${@:2}"; do
+        case "$arg" in
+          -f|--force) force=1 ;;
+          *) branches+=("$arg") ;;
+        esac
+      done
+      if (( ${#branches[@]} == 0 )); then
         print -P "%F{red}error:%f missing <branch> argument"
         return 1
       fi
-      local branches=("${@:2}")
       local root=$(git rev-parse --show-toplevel)
       local wt_list
       wt_list=$(git worktree list --porcelain)
@@ -77,7 +100,9 @@ Manage git worktrees by branch name."
             "$wt_path/bazelw" clean --expunge
           fi
         fi
-        if ! git worktree remove "$wt_path"; then
+        local -a rm_args=()
+        (( force )) && rm_args+=(--force)
+        if ! git worktree remove "${rm_args[@]}" "$wt_path"; then
           failed=1
           continue
         fi