Add option to select base branch for new worktrees

- Thanks to @Juksuu - ThePrimeagen#59

Author: bradsherman

README.md

@@ -107,8 +107,12 @@ Three primary functions should cover your day-to-day.
 The path can be either relative from the git root dir or absoulute path to the worktree.
 
 ```lua
--- Creates a worktree.  Requires the path, branch name, and the upstream
--- Example:
+:-- Creates a worktree.  Requires the path, branch name, and the upstream,
+-- Optionally the base branch from which to create the new worktree and branch can be added
+-- Examples:
+-- Creating new worktree and branch `feat-69` based on develop branch to path `feat-69`
+:lua require("git-worktree").create_worktree("feat-69", "feat-69", "origin", "develop")
+-- Creating new worktree `master` based on master to path `feat-69`
 :lua require("git-worktree").create_worktree("feat-69", "master", "origin")
 
 -- switches to an existing worktree.  Requires the path name

lua/git-worktree/init.lua

@@ -177,7 +177,7 @@ local function change_dirs(path)
     return previous_worktree
 end
 
-local function create_worktree_job(path, branch, found_branch)
+local function create_worktree_job(path, branch, found_branch, base_branch)
     local worktree_add_cmd = "git"
     local worktree_add_args = { "worktree", "add" }
 
@@ -186,10 +186,12 @@ local function create_worktree_job(path, branch, found_branch)
             table.insert(worktree_add_args, path)
             local local_branch_name = branch:gsub("origin/", "")
             table.insert(worktree_add_args, local_branch_name)
+            table.insert(worktree_add_args, base_branch)
         else
             table.insert(worktree_add_args, "-b")
             table.insert(worktree_add_args, branch)
             table.insert(worktree_add_args, path)
+            table.insert(worktree_add_args, base_branch)
         end
     end
 
@@ -308,8 +310,42 @@ local function has_branch(branch, cb)
     end):start()
 end
 
-local function create_worktree(path, branch, upstream, found_branch)
-    local create = create_worktree_job(path, branch, found_branch)
+-- Has branch function to use outside of this file
+-- apparently existing one doesn't work?
+M.has_branch = function(branch)
+    local found = false
+    Job:new({
+        "git",
+        "branch",
+        on_stdout = function(_, data)
+            -- remove marker on current branch
+            data = data:gsub("*", "")
+            data = vim.trim(data)
+            found = found or data == branch
+        end,
+        cwd = git_worktree_root,
+    }):sync()
+
+    return found
+end
+
+local function create_worktree(path, branch, upstream, found_branch, base_branch)
+    has_branch(base_branch, function(found)
+        if not found then
+            status:status("Valid base branch was not defined, using current worktree")
+            base_branch = nil
+        end
+    end)
+
+    local current_branch_job = Job:new({
+        "git",
+        "branch",
+        "--show-current",
+        cwd = vim.loop.cwd(),
+        on_stdout = function(_, data)
+            base_branch = base_branch or data
+        end,
+    })
 
     local worktree_path
     if Path:new(path):is_absolute() then
@@ -364,54 +400,65 @@ local function create_worktree(path, branch, upstream, found_branch)
         end,
     })
 
-    if upstream ~= nil then
-        if M._config.fetch_on_create then
-            create:and_then_on_success(fetch)
-            fetch:and_then_on_success(set_branch)
-        else
-            create:and_then_on_success(set_branch)
-        end
+    current_branch_job:add_on_exit_callback(function()
+        local create = create_worktree_job(path, branch, found_branch, base_branch)
+        if upstream ~= nil then
+            if M._config.fetch_on_create then
+                create:and_then_on_success(fetch)
+                fetch:and_then_on_success(set_branch)
+            else
+                create:and_then_on_success(set_branch)
+            end
 
-        if M._config.autopush then
-            -- These are "optional" operations.
-            -- We have to figure out how we want to handle these...
-            set_branch:and_then(set_push)
-            set_push:and_then(rebase)
-            set_push:after_failure(failure("create_worktree", set_branch.args, worktree_path, true))
-        else
-            set_branch:and_then(rebase)
-        end
+            if M._config.autopush then
+                -- These are "optional" operations.
+                -- We have to figure out how we want to handle these...
+                set_branch:and_then(set_push)
+                set_push:and_then(rebase)
+                set_push:after_failure(failure("create_worktree", set_branch.args, worktree_path, true))
+            else
+                set_branch:and_then(rebase)
+            end
 
-        create:after_failure(failure("create_worktree", create.args, git_worktree_root))
-        if M._config.fetch_on_create then
-            fetch:after_failure(failure("create_worktree", fetch.args, worktree_path))
-        end
+            create:after_failure(failure("create_worktree", create.args, git_worktree_root))
+            if M._config.fetch_on_create then
+                fetch:after_failure(failure("create_worktree", fetch.args, worktree_path))
+            end
 
-        set_branch:after_failure(failure("create_worktree", set_branch.args, worktree_path, true))
+            set_branch:after_failure(failure("create_worktree", set_branch.args, worktree_path, true))
 
-        rebase:after(function()
-            if rebase.code ~= 0 then
-                status:status("Rebase failed, but that's ok.")
-            end
+            rebase:after(function()
+                if rebase.code ~= 0 then
+                    status:status("Rebase failed, but that's ok.")
+                end
 
-            vim.schedule(function()
-                emit_on_change(Enum.Operations.Create, { path = worktree_path, branch = branch, upstream = upstream })
-                M.switch_worktree(worktree_path)
+                vim.schedule(function()
+                    emit_on_change(
+                        Enum.Operations.Create,
+                        { path = worktree_path, branch = branch, upstream = upstream }
+                    )
+                    M.switch_worktree(worktree_path)
+                end)
             end)
-        end)
-    else
-        create:after(function()
-            vim.schedule(function()
-                emit_on_change(Enum.Operations.Create, { path = worktree_path, branch = branch, upstream = upstream })
-                M.switch_worktree(worktree_path)
+        else
+            create:after(function()
+                vim.schedule(function()
+                    emit_on_change(
+                        Enum.Operations.Create,
+                        { path = worktree_path, branch = branch, upstream = upstream }
+                    )
+                    M.switch_worktree(worktree_path)
+                end)
             end)
-        end)
-    end
+        end
+
+        create:start()
+    end)
 
-    create:start()
+    current_branch_job:start()
 end
 
-M.create_worktree = function(path, branch, upstream)
+M.create_worktree = function(path, branch, upstream, base_branch)
     status:reset(8)
 
     if upstream == nil then
@@ -428,7 +475,7 @@ M.create_worktree = function(path, branch, upstream)
         end
 
         has_branch(branch, function(found_branch)
-            create_worktree(path, branch, upstream, found_branch)
+            create_worktree(path, branch, upstream, found_branch, base_branch)
         end)
     end)
 end

lua/telescope/_extensions/git_worktree.lua

@@ -83,7 +83,7 @@ wt_actions.delete_worktree = function(prompt_bufnr)
     end
 end
 
-local create_input_prompt = function(cb)
+local create_input_prompt = function()
     --[[
     local window = Window.centered({
         width = 30,
@@ -103,9 +103,32 @@ local create_input_prompt = function(cb)
     end)
     --]]
     --
+    return vim.fn.input("Path to subtree > ")
+end
+
+local use_current_worktree_as_base_prompt = function()
+    return vim.fn.confirm("Use current worktree as base?", "&Yes\n&No", 1) == 1
+end
+
+local get_base_branch = function(opts, name, branch)
+    local base_branch_selection_opts = opts or {}
+    base_branch_selection_opts.attach_mappings = function()
+        actions.select_default:replace(function(prompt_bufnr, _)
+            local selected_entry = action_state.get_selected_entry()
+            local current_line = action_state.get_current_line()
+
+            actions.close(prompt_bufnr)
+
+            local base_branch = selected_entry ~= nil and selected_entry.value or current_line
+
+            git_worktree.create_worktree(name, branch, nil, base_branch)
+        end)
+
+        -- do we need to replace other default maps?
 
-    local subtree = vim.fn.input("Path to subtree > ")
-    cb(subtree)
+        return true
+    end
+    require("telescope.builtin").git_branches(base_branch_selection_opts)
 end
 
 local pconf = {
@@ -159,8 +182,8 @@ local get_default_opts = function(opts)
 end
 
 local create_worktree = function(opts)
-    opts = get_default_opts(opts)
-    opts.attach_mappings = function()
+    local branch_selection_opts = get_default_opts(opts)
+    branch_selection_opts.attach_mappings = function()
         actions.select_default:replace(function(prompt_bufnr, _)
             local selected_entry = action_state.get_selected_entry()
             local current_line = action_state.get_current_line()
@@ -173,19 +196,29 @@ local create_worktree = function(opts)
                 return
             end
 
-            create_input_prompt(function(name)
-                if name == "" then
-                    name = branch
+            local name = create_input_prompt()
+            if name == "" then
+                name = branch
+            end
+
+            local has_branch = git_worktree.has_branch(branch)
+
+            if not has_branch then
+                if use_current_worktree_as_base_prompt() then
+                    git_worktree.create_worktree(name, branch)
+                else
+                    get_base_branch(opts, name, branch)
                 end
+            else
                 git_worktree.create_worktree(name, branch)
-            end)
+            end
         end)
 
         -- do we need to replace other default maps?
 
         return true
     end
-    require("telescope.builtin").git_branches(opts)
+    require("telescope.builtin").git_branches(branch_selection_opts)
 end
 
 local telescope_git_worktree = function(opts)