adds Tree method to Tree (#224)

This patch adds a new method to the Tree object that allows you to get a child tree from a parent tree by its relative name.

Before this patch, this was only possible with files, using the File method.

The new Tree method has a similar signature to the old File method for consistency.

2 files changed, 57 insertions, 5 deletions

diff --git a/plumbing/object/tree.go b/plumbing/object/tree.go
index 7a8c1a3..f3e03bc 100644
--- a/plumbing/object/tree.go
+++ b/plumbing/object/tree.go
@@ -24,8 +24,9 @@ const (
 
 // New errors defined by this package.
 var (
-	ErrMaxTreeDepth = errors.New("maximum tree depth exceeded")
-	ErrFileNotFound = errors.New("file not found")
+	ErrMaxTreeDepth      = errors.New("maximum tree depth exceeded")
+	ErrFileNotFound      = errors.New("file not found")
+	ErrDirectoryNotFound = errors.New("directory not found")
 )
 
 // Tree is basically like a directory - it references a bunch of other trees
@@ -76,12 +77,31 @@ func (t *Tree) File(path string) (*File, error) {
 
 	blob, err := GetBlob(t.s, e.Hash)
 	if err != nil {
+		if err == plumbing.ErrObjectNotFound {
+			return nil, ErrFileNotFound
+		}
 		return nil, err
 	}
 
 	return NewFile(path, e.Mode, blob), nil
 }
 
+// Tree returns the tree identified by the `path` argument.
+// The path is interpreted as relative to the tree receiver.
+func (t *Tree) Tree(path string) (*Tree, error) {
+	e, err := t.findEntry(path)
+	if err != nil {
+		return nil, ErrDirectoryNotFound
+	}
+
+	tree, err := GetTree(t.s, e.Hash)
+	if err == plumbing.ErrObjectNotFound {
+		return nil, ErrDirectoryNotFound
+	}
+
+	return tree, err
+}
+
 // TreeEntryFile returns the *File for a given *TreeEntry.
 func (t *Tree) TreeEntryFile(e *TreeEntry) (*File, error) {
 	blob, err := GetBlob(t.s, e.Hash)
@@ -106,12 +126,10 @@ func (t *Tree) findEntry(path string) (*TreeEntry, error) {
 	return tree.entry(pathParts[0])
 }
 
-var errDirNotFound = errors.New("directory not found")
-
 func (t *Tree) dir(baseName string) (*Tree, error) {
 	entry, err := t.entry(baseName)
 	if err != nil {
-		return nil, errDirNotFound
+		return nil, ErrDirectoryNotFound
 	}
 
 	obj, err := t.s.EncodedObject(plumbing.TreeObject, entry.Hash)
diff --git a/plumbing/object/tree_test.go b/plumbing/object/tree_test.go
index 00601c1..9b9614b 100644
--- a/plumbing/object/tree_test.go
+++ b/plumbing/object/tree_test.go
@@ -48,6 +48,31 @@ func (s *TreeSuite) TestType(c *C) {
 	c.Assert(s.Tree.Type(), Equals, plumbing.TreeObject)
 }
 
+func (s *TreeSuite) TestTree(c *C) {
+	expectedEntry, ok := s.Tree.m["vendor"]
+	c.Assert(ok, Equals, true)
+	expected := expectedEntry.Hash
+
+	obtainedTree, err := s.Tree.Tree("vendor")
+	c.Assert(err, IsNil)
+	c.Assert(obtainedTree.Hash, Equals, expected)
+}
+
+func (s *TreeSuite) TestTreeNotFound(c *C) {
+	d, err := s.Tree.Tree("not-found")
+	c.Assert(d, IsNil)
+	c.Assert(err, Equals, ErrDirectoryNotFound)
+}
+
+func (s *TreeSuite) TestTreeFailsWithExistingFiles(c *C) {
+	_, err := s.Tree.File("LICENSE")
+	c.Assert(err, IsNil)
+
+	d, err := s.Tree.Tree("LICENSE")
+	c.Assert(d, IsNil)
+	c.Assert(err, Equals, ErrDirectoryNotFound)
+}
+
 func (s *TreeSuite) TestFile(c *C) {
 	f, err := s.Tree.File("LICENSE")
 	c.Assert(err, IsNil)
@@ -60,6 +85,15 @@ func (s *TreeSuite) TestFileNotFound(c *C) {
 	c.Assert(err, Equals, ErrFileNotFound)
 }
 
+func (s *TreeSuite) TestFileFailsWithExistingTrees(c *C) {
+	_, err := s.Tree.Tree("vendor")
+	c.Assert(err, IsNil)
+
+	f, err := s.Tree.File("vendor")
+	c.Assert(f, IsNil)
+	c.Assert(err, Equals, ErrFileNotFound)
+}
+
 func (s *TreeSuite) TestFiles(c *C) {
 	var count int
 	err := s.Tree.Files().ForEach(func(f *File) error {