worktree: checkout, create branch

3 files changed, 116 insertions, 23 deletions

diff --git a/options.go b/options.go
index bbfe244..67a4870 100644
--- a/options.go
+++ b/options.go
@@ -194,13 +194,20 @@ type SubmoduleUpdateOptions struct {
 	RecurseSubmodules SubmoduleRescursivity
 }
 
+var (
+	ErrBranchHashExclusive  = errors.New("Branch and Hash are mutually exclusive")
+	ErrCreateRequiresBranch = errors.New("Branch is mandatory when Create is used")
+)
+
 // CheckoutOptions describes how a checkout 31operation should be performed.
 type CheckoutOptions struct {
 	// Hash to be checked out, if used HEAD will in detached mode. Branch and
-	// Hash are mutual exclusive.
+	// Hash are mutually exclusive, if Create is not used.
 	Hash plumbing.Hash
 	// Branch to be checked out, if Branch and Hash are empty is set to `master`.
 	Branch plumbing.ReferenceName
+	// Create a new branch named Branch and start it at Hash.
+	Create bool
 	// Force, if true when switching branches, proceed even if the index or the
 	// working tree differs from HEAD. This is used to throw away local changes
 	Force bool
@@ -208,6 +215,14 @@ type CheckoutOptions struct {
 
 // Validate validates the fields and sets the default values.
 func (o *CheckoutOptions) Validate() error {
+	if !o.Create && !o.Hash.IsZero() && o.Branch != "" {
+		return ErrBranchHashExclusive
+	}
+
+	if o.Create && o.Branch == "" {
+		return ErrCreateRequiresBranch
+	}
+
 	if o.Branch == "" {
 		o.Branch = plumbing.Master
 	}
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
diff --git a/worktree_test.go b/worktree_test.go
index 4c9907b..a6c7b06 100644
--- a/worktree_test.go
+++ b/worktree_test.go
@@ -173,47 +173,96 @@ func (s *WorktreeSuite) TestCheckoutIndexOS(c *C) {
 	c.Assert(idx.Entries[0].GID, Not(Equals), uint32(0))
 }
 
-func (s *WorktreeSuite) TestCheckoutChange(c *C) {
-	fs := memfs.New()
+func (s *WorktreeSuite) TestCheckoutBranch(c *C) {
 	w := &Worktree{
 		r:  s.Repository,
-		fs: fs,
+		fs: memfs.New(),
 	}
 
-	err := w.Checkout(&CheckoutOptions{})
+	err := w.Checkout(&CheckoutOptions{
+		Branch: "refs/heads/branch",
+	})
 	c.Assert(err, IsNil)
 
 	head, err := w.r.Head()
 	c.Assert(err, IsNil)
-	c.Assert(head.Name().String(), Equals, "refs/heads/master")
+	c.Assert(head.Name().String(), Equals, "refs/heads/branch")
 
 	status, err := w.Status()
 	c.Assert(err, IsNil)
 	c.Assert(status.IsClean(), Equals, true)
+}
 
-	_, err = fs.Stat("README")
-	c.Assert(err, Equals, os.ErrNotExist)
-	_, err = fs.Stat("vendor")
-	c.Assert(err, Equals, nil)
+func (s *WorktreeSuite) TestCheckoutCreateWithHash(c *C) {
+	w := &Worktree{
+		r:  s.Repository,
+		fs: memfs.New(),
+	}
 
-	err = w.Checkout(&CheckoutOptions{
-		Branch: "refs/heads/branch",
+	err := w.Checkout(&CheckoutOptions{
+		Create: true,
+		Branch: "refs/heads/foo",
+		Hash:   plumbing.NewHash("35e85108805c84807bc66a02d91535e1e24b38b9"),
 	})
 	c.Assert(err, IsNil)
 
-	status, err = w.Status()
+	head, err := w.r.Head()
+	c.Assert(err, IsNil)
+	c.Assert(head.Name().String(), Equals, "refs/heads/foo")
+	c.Assert(head.Hash(), Equals, plumbing.NewHash("35e85108805c84807bc66a02d91535e1e24b38b9"))
+
+	status, err := w.Status()
 	c.Assert(err, IsNil)
 	c.Assert(status.IsClean(), Equals, true)
+}
 
-	_, err = fs.Stat("README")
-	c.Assert(err, Equals, nil)
+func (s *WorktreeSuite) TestCheckoutCreate(c *C) {
+	w := &Worktree{
+		r:  s.Repository,
+		fs: memfs.New(),
+	}
 
-	_, err = fs.Stat("vendor")
-	c.Assert(err, Equals, os.ErrNotExist)
+	err := w.Checkout(&CheckoutOptions{
+		Create: true,
+		Branch: "refs/heads/foo",
+	})
+	c.Assert(err, IsNil)
 
-	head, err = w.r.Head()
+	head, err := w.r.Head()
 	c.Assert(err, IsNil)
-	c.Assert(head.Name().String(), Equals, "refs/heads/branch")
+	c.Assert(head.Name().String(), Equals, "refs/heads/foo")
+	c.Assert(head.Hash(), Equals, plumbing.NewHash("6ecf0ef2c2dffb796033e5a02219af86ec6584e5"))
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status.IsClean(), Equals, true)
+}
+
+func (s *WorktreeSuite) TestCheckoutBranchAndHash(c *C) {
+	w := &Worktree{
+		r:  s.Repository,
+		fs: memfs.New(),
+	}
+
+	err := w.Checkout(&CheckoutOptions{
+		Branch: "refs/heads/foo",
+		Hash:   plumbing.NewHash("35e85108805c84807bc66a02d91535e1e24b38b9"),
+	})
+
+	c.Assert(err, Equals, ErrBranchHashExclusive)
+}
+
+func (s *WorktreeSuite) TestCheckoutCreateMissingBranch(c *C) {
+	w := &Worktree{
+		r:  s.Repository,
+		fs: memfs.New(),
+	}
+
+	err := w.Checkout(&CheckoutOptions{
+		Create: true,
+	})
+
+	c.Assert(err, Equals, ErrCreateRequiresBranch)
 }
 
 func (s *WorktreeSuite) TestCheckoutTag(c *C) {