new methods Worktree.[AddGlob|RemoveBlob] and recursive Worktree.[Add|Remove]

Signed-off-by: Máximo Cuadros <mcuadros@gmail.com>

2 files changed, 254 insertions, 77 deletions

diff --git a/worktree_status.go b/worktree_status.go
index 4141381..2cac78e 100644
--- a/worktree_status.go
+++ b/worktree_status.go
@@ -5,6 +5,8 @@ import (
 	"errors"
 	"io"
 	"os"
+	"path"
+	"path/filepath"
 
 	"gopkg.in/src-d/go-billy.v4/util"
 	"gopkg.in/src-d/go-git.v4/plumbing"
@@ -24,10 +26,8 @@ var (
 	// the worktree.
 	ErrDestinationExists = errors.New("destination exists")
 	// ErrGlobNoMatches in an AddGlob if the glob pattern does not match any
-	// file in the worktree.ErrNotDirectory
+	// files in the worktree.
 	ErrGlobNoMatches = errors.New("glob pattern did not match any files")
-	// ErrNotDirectory  in an AddDirectory if the path is not a directory.
-	ErrNotDirectory = errors.New("path is not a directory")
 )
 
 // Status returns the working tree status.
@@ -252,8 +252,12 @@ func diffTreeIsEquals(a, b noder.Hasher) bool {
 
 // Add adds the file contents of a file in the worktree to the index. if the
 // file is already staged in the index no error is returned. If a file deleted
-// from the Workspace is given, the file is removed from the index.
+// from the Workspace is given, the file is removed from the index. If a
+// directory given, adds the files and all his sub-directories recursively in
+// the worktree to the index. If any of the files is already staged in the index
+// no error is returned. When path is a file, the blob.Hash is returned.
 func (w *Worktree) Add(path string) (plumbing.Hash, error) {
+	// TODO(mcuadros): remove plumbing.Hash from signature at v5.
 	s, err := w.Status()
 	if err != nil {
 		return plumbing.ZeroHash, err
@@ -264,67 +268,41 @@ func (w *Worktree) Add(path string) (plumbing.Hash, error) {
 		return plumbing.ZeroHash, err
 	}
 
-	added, h, err := w.doAdd(idx, s, path)
-	if err != nil {
-		return h, err
-	}
-
-	if !added {
-		return h, nil
-	}
+	var h plumbing.Hash
+	var added bool
 
-	return h, w.r.Storer.SetIndex(idx)
-}
-
-// AddDirectory adds the files contents of a directory and all his
-// sub-directories recursively in the worktree to the index. If any of the
-// file is already staged in the index no error is returned.
-func (w *Worktree) AddDirectory(path string) error {
 	fi, err := w.Filesystem.Lstat(path)
-	if err != nil {
-		return err
-	}
-
-	if !fi.IsDir() {
-		return ErrNotDirectory
-	}
-
-	s, err := w.Status()
-	if err != nil {
-		return err
-	}
-
-	idx, err := w.r.Storer.Index()
-	if err != nil {
-		return err
+	if err != nil || !fi.IsDir() {
+		added, h, err = w.doAddFile(idx, s, path)
+	} else {
+		added, err = w.doAddDirectory(idx, s, path)
 	}
 
-	added, err := w.doAddDirectory(idx, s, path)
 	if err != nil {
-		return err
+		return h, err
 	}
 
 	if !added {
-		return nil
+		return h, nil
 	}
 
-	return w.r.Storer.SetIndex(idx)
+	return h, w.r.Storer.SetIndex(idx)
 }
 
-func (w *Worktree) doAddDirectory(idx *index.Index, s Status, path string) (added bool, err error) {
-	files, err := w.Filesystem.ReadDir(path)
+func (w *Worktree) doAddDirectory(idx *index.Index, s Status, directory string) (added bool, err error) {
+	files, err := w.Filesystem.ReadDir(directory)
 	if err != nil {
 		return false, err
 	}
 
 	for _, file := range files {
-		name := w.Filesystem.Join(path, file.Name())
+		name := path.Join(directory, file.Name())
 
 		var a bool
 		if file.IsDir() {
 			a, err = w.doAddDirectory(idx, s, name)
 		} else {
-			a, _, err = w.doAdd(idx, s, name)
+			a, _, err = w.doAddFile(idx, s, name)
 		}
 
 		if err != nil {
@@ -337,12 +315,11 @@ func (w *Worktree) doAddDirectory(idx *index.Index, s Status, path string) (adde
 	}
 
 	return
-
 }
 
-// AddGlob given a glob pattern adds all the matching files content and all his
-// sub-directories recursively in the worktree to the index. If any of the
-// file is already staged in the index no error is returned.
+// AddGlob adds all paths, matching pattern, to the index. If pattern matches a
+// directory path, all directory contents are added to the index recursively. No
+// error is returned if all matching paths are already staged in index.
 func (w *Worktree) AddGlob(pattern string) error {
 	files, err := util.Glob(w.Filesystem, pattern)
 	if err != nil {
@@ -374,7 +351,7 @@ func (w *Worktree) AddGlob(pattern string) error {
 		if fi.IsDir() {
 			added, err = w.doAddDirectory(idx, s, file)
 		} else {
-			added, _, err = w.doAdd(idx, s, file)
+			added, _, err = w.doAddFile(idx, s, file)
 		}
 
 		if err != nil {
@@ -393,9 +370,13 @@ func (w *Worktree) AddGlob(pattern string) error {
 	return nil
 }
 
-// doAdd create a new blob from path and update the index, added is true if
+// doAddFile create a new blob from path and update the index, added is true if
 // the file added is different from the index.
-func (w *Worktree) doAdd(idx *index.Index, s Status, path string) (added bool, h plumbing.Hash, err error) {
+func (w *Worktree) doAddFile(idx *index.Index, s Status, path string) (added bool, h plumbing.Hash, err error) {
+	if s.File(path).Worktree == Unmodified {
+		return false, h, nil
+	}
+
 	h, err = w.copyFileToStorage(path)
 	if err != nil {
 		if os.IsNotExist(err) {
@@ -406,10 +387,6 @@ func (w *Worktree) doAdd(idx *index.Index, s Status, path string) (added bool, h
 		return
 	}
 
-	if s.File(path).Worktree == Unmodified {
-		return false, h, nil
-	}
-
 	if err := w.addOrUpdateFileToIndex(idx, path, h); err != nil {
 		return false, h, err
 	}
@@ -486,10 +463,7 @@ func (w *Worktree) addOrUpdateFileToIndex(idx *index.Index, filename string, h p
 }
 
 func (w *Worktree) doAddFileToIndex(idx *index.Index, filename string, h plumbing.Hash) error {
-	e := &index.Entry{Name: filename}
-	idx.Entries = append(idx.Entries, e)
-
-	return w.doUpdateFileToIndex(e, filename, h)
+	return w.doUpdateFileToIndex(idx.Add(filename), filename, h)
 }
 
 func (w *Worktree) doUpdateFileToIndex(e *index.Entry, filename string, h plumbing.Hash) error {
@@ -515,21 +489,79 @@ func (w *Worktree) doUpdateFileToIndex(e *index.Entry, filename string, h plumbi
 
 // Remove removes files from the working tree and from the index.
 func (w *Worktree) Remove(path string) (plumbing.Hash, error) {
+	// TODO(mcuadros): remove plumbing.Hash from signature at v5.
 	idx, err := w.r.Storer.Index()
 	if err != nil {
 		return plumbing.ZeroHash, err
 	}
 
-	hash, err := w.deleteFromIndex(idx, path)
+	var h plumbing.Hash
+
+	fi, err := w.Filesystem.Lstat(path)
+	if err != nil || !fi.IsDir() {
+		h, err = w.doRemoveFile(idx, path)
+	} else {
+		_, err = w.doRemoveDirectory(idx, path)
+	}
 	if err != nil {
-		return plumbing.ZeroHash, err
+		return h, err
 	}
 
-	if err := w.deleteFromFilesystem(path); err != nil {
-		return hash, err
+	return h, w.r.Storer.SetIndex(idx)
+}
+
+func (w *Worktree) doRemoveDirectory(idx *index.Index, directory string) (removed bool, err error) {
+	files, err := w.Filesystem.ReadDir(directory)
+	if err != nil {
+		return false, err
 	}
 
-	return hash, w.r.Storer.SetIndex(idx)
+	for _, file := range files {
+		name := path.Join(directory, file.Name())
+
+		var r bool
+		if file.IsDir() {
+			r, err = w.doRemoveDirectory(idx, name)
+		} else {
+			_, err = w.doRemoveFile(idx, name)
+			if err == index.ErrEntryNotFound {
+				err = nil
+			}
+		}
+
+		if err != nil {
+			return
+		}
+
+		if !removed && r {
+			removed = true
+		}
+	}
+
+	err = w.removeEmptyDirectory(directory)
+	return
+}
+
+func (w *Worktree) removeEmptyDirectory(path string) error {
+	files, err := w.Filesystem.ReadDir(path)
+	if err != nil {
+		return err
+	}
+
+	if len(files) != 0 {
+		return nil
+	}
+
+	return w.Filesystem.Remove(path)
+}
+
+func (w *Worktree) doRemoveFile(idx *index.Index, path string) (plumbing.Hash, error) {
+	hash, err := w.deleteFromIndex(idx, path)
+	if err != nil {
+		return plumbing.ZeroHash, err
+	}
+
+	return hash, w.deleteFromFilesystem(path)
 }
 
 func (w *Worktree) deleteFromIndex(idx *index.Index, path string) (plumbing.Hash, error) {
@@ -550,9 +582,43 @@ func (w *Worktree) deleteFromFilesystem(path string) error {
 	return err
 }
 
+// RemoveGlob removes all paths, matching pattern, from the index. If pattern
+// matches a directory path, all directory contents are removed from the index
+// recursively.
+func (w *Worktree) RemoveGlob(pattern string) error {
+	idx, err := w.r.Storer.Index()
+	if err != nil {
+		return err
+	}
+
+	entries, err := idx.Glob(pattern)
+	if err != nil {
+		return err
+	}
+
+	for _, e := range entries {
+		file := filepath.FromSlash(e.Name)
+		if _, err := w.Filesystem.Lstat(file); err != nil && !os.IsNotExist(err) {
+			return err
+		}
+
+		if _, err := w.doRemoveFile(idx, file); err != nil {
+			return err
+		}
+
+		dir, _ := filepath.Split(file)
+		if err := w.removeEmptyDirectory(dir); err != nil {
+			return err
+		}
+	}
+
+	return w.r.Storer.SetIndex(idx)
+}
+
 // Move moves or rename a file in the worktree and the index, directories are
 // not supported.
 func (w *Worktree) Move(from, to string) (plumbing.Hash, error) {
+	// TODO(mcuadros): support directories and/or implement support for glob
 	if _, err := w.Filesystem.Lstat(from); err != nil {
 		return plumbing.ZeroHash, err
 	}
diff --git a/worktree_test.go b/worktree_test.go
index 8f30e82..cb2e5e2 100644
--- a/worktree_test.go
+++ b/worktree_test.go
@@ -1196,8 +1196,9 @@ func (s *WorktreeSuite) TestAddDirectory(c *C) {
 	err = util.WriteFile(w.Filesystem, "qux/baz/bar", []byte("BAR"), 0755)
 	c.Assert(err, IsNil)
 
-	err = w.AddDirectory("qux")
+	h, err := w.Add("qux")
 	c.Assert(err, IsNil)
+	c.Assert(h.IsZero(), Equals, true)
 
 	idx, err = w.r.Storer.Index()
 	c.Assert(err, IsNil)
@@ -1224,23 +1225,13 @@ func (s *WorktreeSuite) TestAddDirectory(c *C) {
 	c.Assert(file.Worktree, Equals, Unmodified)
 }
 
-func (s *WorktreeSuite) TestAddDirectoryErrorNotDirectory(c *C) {
-	r, _ := Init(memory.NewStorage(), memfs.New())
-	w, _ := r.Worktree()
-
-	err := util.WriteFile(w.Filesystem, "foo", []byte("FOO"), 0755)
-	c.Assert(err, IsNil)
-
-	err = w.AddDirectory("foo")
-	c.Assert(err, Equals, ErrNotDirectory)
-}
-
 func (s *WorktreeSuite) TestAddDirectoryErrorNotFound(c *C) {
 	r, _ := Init(memory.NewStorage(), memfs.New())
 	w, _ := r.Worktree()
 
-	err := w.AddDirectory("foo")
+	h, err := w.Add("foo")
 	c.Assert(err, NotNil)
+	c.Assert(h.IsZero(), Equals, true)
 }
 
 func (s *WorktreeSuite) TestAddGlob(c *C) {
@@ -1264,7 +1255,7 @@ func (s *WorktreeSuite) TestAddGlob(c *C) {
 	err = util.WriteFile(w.Filesystem, "qux/bar/baz", []byte("BAZ"), 0755)
 	c.Assert(err, IsNil)
 
-	err = w.AddGlob("qux/b*")
+	err = w.AddGlob(w.Filesystem.Join("qux", "b*"))
 	c.Assert(err, IsNil)
 
 	idx, err = w.r.Storer.Index()
@@ -1339,6 +1330,58 @@ func (s *WorktreeSuite) TestRemoveNotExistentEntry(c *C) {
 	c.Assert(err, NotNil)
 }
 
+func (s *WorktreeSuite) TestRemoveDirectory(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:          s.Repository,
+		Filesystem: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	hash, err := w.Remove("json")
+	c.Assert(hash.IsZero(), Equals, true)
+	c.Assert(err, IsNil)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 2)
+	c.Assert(status.File("json/long.json").Staging, Equals, Deleted)
+	c.Assert(status.File("json/short.json").Staging, Equals, Deleted)
+
+	_, err = w.Filesystem.Stat("json")
+	c.Assert(os.IsNotExist(err), Equals, true)
+}
+
+func (s *WorktreeSuite) TestRemoveDirectoryUntracked(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:          s.Repository,
+		Filesystem: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	err = util.WriteFile(w.Filesystem, "json/foo", []byte("FOO"), 0755)
+	c.Assert(err, IsNil)
+
+	hash, err := w.Remove("json")
+	c.Assert(hash.IsZero(), Equals, true)
+	c.Assert(err, IsNil)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 3)
+	c.Assert(status.File("json/long.json").Staging, Equals, Deleted)
+	c.Assert(status.File("json/short.json").Staging, Equals, Deleted)
+	c.Assert(status.File("json/foo").Staging, Equals, Untracked)
+
+	_, err = w.Filesystem.Stat("json")
+	c.Assert(err, IsNil)
+}
+
 func (s *WorktreeSuite) TestRemoveDeletedFromWorktree(c *C) {
 	fs := memfs.New()
 	w := &Worktree{
@@ -1362,6 +1405,74 @@ func (s *WorktreeSuite) TestRemoveDeletedFromWorktree(c *C) {
 	c.Assert(status.File("LICENSE").Staging, Equals, Deleted)
 }
 
+func (s *WorktreeSuite) TestRemoveGlob(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:          s.Repository,
+		Filesystem: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	err = w.RemoveGlob(w.Filesystem.Join("json", "l*"))
+	c.Assert(err, IsNil)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 1)
+	c.Assert(status.File("json/long.json").Staging, Equals, Deleted)
+}
+
+func (s *WorktreeSuite) TestRemoveGlobDirectory(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:          s.Repository,
+		Filesystem: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	err = w.RemoveGlob("js*")
+	c.Assert(err, IsNil)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 2)
+	c.Assert(status.File("json/short.json").Staging, Equals, Deleted)
+	c.Assert(status.File("json/long.json").Staging, Equals, Deleted)
+
+	_, err = w.Filesystem.Stat("json")
+	c.Assert(os.IsNotExist(err), Equals, true)
+}
+
+func (s *WorktreeSuite) TestRemoveGlobDirectoryDeleted(c *C) {
+	fs := memfs.New()
+	w := &Worktree{
+		r:          s.Repository,
+		Filesystem: fs,
+	}
+
+	err := w.Checkout(&CheckoutOptions{Force: true})
+	c.Assert(err, IsNil)
+
+	err = fs.Remove("json/short.json")
+	c.Assert(err, IsNil)
+
+	err = util.WriteFile(w.Filesystem, "json/foo", []byte("FOO"), 0755)
+	c.Assert(err, IsNil)
+
+	err = w.RemoveGlob("js*")
+	c.Assert(err, IsNil)
+
+	status, err := w.Status()
+	c.Assert(err, IsNil)
+	c.Assert(status, HasLen, 3)
+	c.Assert(status.File("json/short.json").Staging, Equals, Deleted)
+	c.Assert(status.File("json/long.json").Staging, Equals, Deleted)
+}
+
 func (s *WorktreeSuite) TestMove(c *C) {
 	fs := memfs.New()
 	w := &Worktree{