Implement dict protocol for trees.

It is rather intuitive to consider trees as a dict of objects (like
a directory could be seen as a dict of files).

Author: nud

lib/git/tree.py

@@ -15,7 +15,7 @@ def __init__(self, repo, **kwargs):
         self.id = None
         self.mode = None
         self.name = None
-        self.contents = None
+        self._contents = None
 
         for k, v in kwargs.items():
             setattr(self, k, v)
@@ -27,11 +27,11 @@ def __bake__(self):
             treeish = treeish + ':'
 
         # Read the tree contents.
-        self.contents = {}
+        self._contents = {}
         for line in self.repo.git.ls_tree(self.id).splitlines():
             obj = self.content_from_string(self.repo, line)
-            if obj:
-                self.contents[obj.name] = obj
+            if obj is not None:
+                self._contents[obj.name] = obj
 
     def content_from_string(self, repo, text):
         """
@@ -74,11 +74,37 @@ def __div__(self, file):
         Returns
             ``GitPython.Blob`` or ``GitPython.Tree`` or ``None`` if not found
         """
-        return self.contents.get(file)
+        return self.get(file)
 
     @property
     def basename(self):
         os.path.basename(self.name)
 
     def __repr__(self):
         return '<GitPython.Tree "%s">' % self.id
+
+    # Implement the basics of the dict protocol:
+    # directories/trees can be seen as object dicts.
+    def __getitem__(self, key):
+        return self._contents[key]
+
+    def __iter__(self):
+        return iter(self._contents)
+
+    def __len__(self, keys):
+        return len(self._contents)
+
+    def __contains__(self, key):
+        return key in self._contents
+
+    def get(self, key):
+        return self._contents.get(key)
+
+    def items(self):
+        return self._contents.items()
+
+    def keys(self):
+        return self._contents.keys()
+
+    def values(self):
+        return self._contents.values()

test/git/test_repo.py

@@ -96,8 +96,8 @@ def test_tree(self, git):
 
         tree = self.repo.tree('master')
 
-        assert_equal(4, len([c for c in tree.contents.values() if isinstance(c, Blob)]))
-        assert_equal(3, len([c for c in tree.contents.values() if isinstance(c, Tree)]))
+        assert_equal(4, len([c for c in tree.values() if isinstance(c, Blob)]))
+        assert_equal(3, len([c for c in tree.values() if isinstance(c, Tree)]))
 
         assert_true(git.called)
         assert_equal(git.call_args, (('ls_tree', 'master'), {}))

test/git/test_tree.py

@@ -18,9 +18,9 @@ def test_contents_should_cache(self, git):
 
         tree = self.repo.tree('master')
 
-        child = tree.contents['grit']
-        child.contents
-        child.contents
+        child = tree['grit']
+        child.items()
+        child.items()
 
         assert_true(git.called)
         assert_equal(2, git.call_count)
@@ -96,6 +96,55 @@ def test_slash_with_commits(self, git):
         assert_true(git.called)
         assert_equal(git.call_args, (('ls_tree', 'master'), {}))
 
+    @patch(Blob, 'size')
+    @patch(Git, '_call_process')
+    def test_dict(self, blob, git):
+        git.return_value = fixture('ls_tree_a')
+        blob.return_value = 1
+
+        tree = self.repo.tree('master')
+
+        assert_equal('aa06ba24b4e3f463b3c4a85469d0fb9e5b421cf8', tree['lib'].id)
+        assert_equal('8b1e02c0fb554eed2ce2ef737a68bb369d7527df', tree['README.txt'].id)
+
+        assert_true(git.called)
+        assert_equal(git.call_args, (('ls_tree', 'master'), {}))
+
+    @patch(Blob, 'size')
+    @patch(Git, '_call_process')
+    def test_dict_with_zero_length_file(self, blob, git):
+        git.return_value = fixture('ls_tree_a')
+        blob.return_value = 0
+
+        tree = self.repo.tree('master')
+
+        assert_not_none(tree['README.txt'])
+        assert_equal('8b1e02c0fb554eed2ce2ef737a68bb369d7527df', tree['README.txt'].id)
+
+        assert_true(git.called)
+        assert_equal(git.call_args, (('ls_tree', 'master'), {}))
+
+    @patch(Git, '_call_process')
+    def test_dict_with_commits(self, git):
+        git.return_value = fixture('ls_tree_commit')
+
+        tree = self.repo.tree('master')
+
+        assert_none(tree.get('bar'))
+        assert_equal('2afb47bcedf21663580d5e6d2f406f08f3f65f19', tree['foo'].id)
+        assert_equal('f623ee576a09ca491c4a27e48c0dfe04be5f4a2e', tree['baz'].id)
+
+        assert_true(git.called)
+        assert_equal(git.call_args, (('ls_tree', 'master'), {}))
+
+    @patch(Git, '_call_process')
+    @raises(KeyError)
+    def test_dict_with_non_existant_file(self, git):
+        git.return_value = fixture('ls_tree_commit')
+
+        tree = self.repo.tree('master')
+        tree['bar']
+
     def test_repr(self):
         self.tree = Tree(self.repo, **{'id': 'abc'})
         assert_equal('<GitPython.Tree "abc">', repr(self.tree))