From 9afc47237c301ecee66619d1ef8ec286185cb070 Mon Sep 17 00:00:00 2001
From: Máximo Cuadros <mcuadros@gmail.com>
Date: Tue, 18 Jul 2017 21:41:34 +0200
Subject: worktree: checkout, create branch

---
 worktree.go | 37 +++++++++++++++++++++++++++++++++----
 1 file changed, 33 insertions(+), 4 deletions(-)

(limited to 'worktree.go')

diff --git a/worktree.go b/worktree.go
index ae1ab28..e0f5fdf 100644
--- a/worktree.go
+++ b/worktree.go
@@ -38,8 +38,14 @@ func (w *Worktree) Checkout(opts *CheckoutOptions) error {
 		return err
 	}
 
+	if opts.Create {
+		if err := w.createBranch(opts); err != nil {
+			return err
+		}
+	}
+
 	if !opts.Force {
-		unstaged, err := w.cointainsUnstagedChanges()
+		unstaged, err := w.containsUnstagedChanges()
 		if err != nil {
 			return err
 		}
@@ -59,7 +65,7 @@ func (w *Worktree) Checkout(opts *CheckoutOptions) error {
 		ro.Mode = HardReset
 	}
 
-	if !opts.Hash.IsZero() {
+	if !opts.Hash.IsZero() && !opts.Create {
 		err = w.setHEADToCommit(opts.Hash)
 	} else {
 		err = w.setHEADToBranch(opts.Branch, c)
@@ -71,6 +77,29 @@ func (w *Worktree) Checkout(opts *CheckoutOptions) error {
 
 	return w.Reset(ro)
 }
+func (w *Worktree) createBranch(opts *CheckoutOptions) error {
+	_, err := w.r.Storer.Reference(opts.Branch)
+	if err == nil {
+		return fmt.Errorf("a branch named %q already exists", opts.Branch)
+	}
+
+	if err != plumbing.ErrReferenceNotFound {
+		return err
+	}
+
+	if opts.Hash.IsZero() {
+		ref, err := w.r.Head()
+		if err != nil {
+			return err
+		}
+
+		opts.Hash = ref.Hash()
+	}
+
+	return w.r.Storer.SetReference(
+		plumbing.NewHashReference(opts.Branch, opts.Hash),
+	)
+}
 
 func (w *Worktree) getCommitFromCheckoutOptions(opts *CheckoutOptions) (plumbing.Hash, error) {
 	if !opts.Hash.IsZero() {
@@ -133,7 +162,7 @@ func (w *Worktree) Reset(opts *ResetOptions) error {
 	}
 
 	if opts.Mode == MergeReset {
-		unstaged, err := w.cointainsUnstagedChanges()
+		unstaged, err := w.containsUnstagedChanges()
 		if err != nil {
 			return err
 		}
@@ -171,7 +200,7 @@ func (w *Worktree) Reset(opts *ResetOptions) error {
 	return w.setHEADCommit(opts.Commit)
 }
 
-func (w *Worktree) cointainsUnstagedChanges() (bool, error) {
+func (w *Worktree) containsUnstagedChanges() (bool, error) {
 	ch, err := w.diffStagingWithWorktree()
 	if err != nil {
 		return false, err
-- 
cgit