recursiveloader: Support getting dict of all file entries

2 files changed, 60 insertions, 0 deletions

diff --git a/gemato/recursiveloader.py b/gemato/recursiveloader.py
index c996944..6cc01b4 100644
--- a/gemato/recursiveloader.py
+++ b/gemato/recursiveloader.py
@@ -153,3 +153,24 @@ class ManifestRecursiveLoader(object):
                     if e.path == filename:
                         return e
         return None
+
+    def get_file_entry_dict(self, path=''):
+        """
+        Find all file entries that apply to paths starting with @path.
+        Return a dictionary mapping relative paths to entries. Raises
+        an exception if multiple entries for file collide.
+        """
+
+        self.load_manifests_for_path(path, recursive=True)
+        out = {}
+        for relpath, m in self._iter_manifests_for_path(path, recursive=True):
+            for e in m.entries:
+                if isinstance(e, gemato.manifest.ManifestEntryDIST):
+                    # distfiles are not local files, so skip them
+                    pass
+                elif isinstance(e, gemato.manifest.ManifestPathEntry):
+                    fullpath = os.path.join(relpath, e.path)
+                    if gemato.util.path_starts_with(fullpath, path):
+                        # TODO: implement conflict detection
+                        out[fullpath] = e
+        return out
diff --git a/tests/test_recursiveloader.py b/tests/test_recursiveloader.py
index d9dfa83..26c7259 100644
--- a/tests/test_recursiveloader.py
+++ b/tests/test_recursiveloader.py
@@ -175,6 +175,45 @@ DATA test 0 MD5 d41d8cd98f00b204e9800998ecf8427e
         self.assertRaises(gemato.verify.ManifestMismatch,
                 m.assert_path_verifies, 'sub/stray')
 
+    def test_get_file_entry_dict(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+            os.path.join(self.dir, 'Manifest'))
+        entries = m.get_file_entry_dict('')
+        self.assertSetEqual(frozenset(entries),
+            frozenset((
+                'other/Manifest',
+                'sub/Manifest',
+                'sub/nonstray',
+                'sub/deeper/Manifest',
+                'sub/deeper/test',
+            )))
+        self.assertEqual(entries['other/Manifest'].path, 'other/Manifest')
+        self.assertEqual(entries['sub/Manifest'].path, 'sub/Manifest')
+        self.assertEqual(entries['sub/nonstray'].path, 'nonstray')
+        self.assertEqual(entries['sub/deeper/Manifest'].path, 'deeper/Manifest')
+        self.assertEqual(entries['sub/deeper/test'].path, 'test')
+
+    def test_get_file_entry_dict_for_sub(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+            os.path.join(self.dir, 'Manifest'))
+        entries = m.get_file_entry_dict('sub')
+        self.assertSetEqual(frozenset(entries),
+            frozenset((
+                'sub/Manifest',
+                'sub/nonstray',
+                'sub/deeper/Manifest',
+                'sub/deeper/test',
+            )))
+        self.assertEqual(entries['sub/Manifest'].path, 'sub/Manifest')
+        self.assertEqual(entries['sub/nonstray'].path, 'nonstray')
+        self.assertEqual(entries['sub/deeper/Manifest'].path, 'deeper/Manifest')
+        self.assertEqual(entries['sub/deeper/test'].path, 'test')
+
+    def test_get_file_entry_dict_for_invalid(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+            os.path.join(self.dir, 'Manifest'))
+        self.assertDictEqual(m.get_file_entry_dict('nonexist'), {})
+
 
 class MultipleManifestTest(unittest.TestCase):
     DIRS = ['sub']