recursiveloader: Support compressed Manifests

2 files changed, 81 insertions, 4 deletions

diff --git a/gemato/recursiveloader.py b/gemato/recursiveloader.py
index 1ea067e..1b4fd14 100644
--- a/gemato/recursiveloader.py
+++ b/gemato/recursiveloader.py
@@ -3,7 +3,6 @@
 # (c) 2017 Michał Górny
 # Licensed under the terms of 2-clause BSD license
 
-import io
 import os.path
 
 import gemato.exceptions
@@ -40,13 +39,15 @@ class ManifestRecursiveLoader(object):
         """
         Load a single Manifest file whose relative path within Manifest
         tree is @relpath. If @verify_entry is not null, the Manifest
-        file is verified against the entry.
+        file is verified against the entry. If the file is compressed,
+        it is decompressed transparently.
         """
         m = gemato.manifest.ManifestFile()
         path = os.path.join(self.root_directory, relpath)
         if verify_entry is not None:
             gemato.verify.assert_path_verifies(path, verify_entry)
-        with io.open(path, 'r', encoding='utf8') as f:
+        with gemato.compression.open_potentially_compressed_path(
+                path, 'r', encoding='utf8') as f:
             m.load(f)
             st = os.fstat(f.fileno())
             self.manifest_device = st.st_dev
@@ -264,7 +265,8 @@ class ManifestRecursiveLoader(object):
                 fpath = os.path.join(relpath, f)
                 # skip top-level Manifest, we obviously can't have
                 # an entry for it
-                if fpath == 'Manifest':
+                if fpath in (gemato.compression
+                        .get_potential_compressed_names('Manifest')):
                     continue
                 fe = entry_dict.pop(fpath, None)
                 self._verify_one_file(os.path.join(dirpath, f), fe, strict)
diff --git a/tests/test_recursiveloader.py b/tests/test_recursiveloader.py
index bc98fb7..92aed2d 100644
--- a/tests/test_recursiveloader.py
+++ b/tests/test_recursiveloader.py
@@ -3,7 +3,10 @@
 # (c) 2017 Michał Górny
 # Licensed under the terms of 2-clause BSD license
 
+import base64
 import datetime
+import gzip
+import io
 import os
 
 import gemato.exceptions
@@ -844,3 +847,75 @@ class UnreadableDirectoryTest(TempDirTestCase):
         m = gemato.recursiveloader.ManifestRecursiveLoader(
             os.path.join(self.dir, 'Manifest'))
         self.assertRaises(OSError, m.assert_directory_verifies)
+
+
+class CompressedTopManifestTest(TempDirTestCase):
+    """
+    Test a tree with top-level Manifest being compressed.
+    """
+
+    MANIFEST = b'''
+DATA test 0 MD5 d41d8cd98f00b204e9800998ecf8427e
+'''
+    FILES = {
+        'test': u'',
+    }
+
+    def setUp(self):
+        super(CompressedTopManifestTest, self).setUp()
+        self.manifest_gz = os.path.join(self.dir, 'Manifest.gz')
+        with gzip.GzipFile(self.manifest_gz, 'wb') as f:
+            f.write(self.MANIFEST)
+
+    def tearDown(self):
+        os.unlink(self.manifest_gz)
+        super(CompressedTopManifestTest, self).tearDown()
+
+    def test_find_path_entry(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+                self.manifest_gz)
+        self.assertEqual(m.find_path_entry('test').path, 'test')
+
+    def test_assert_directory_verifies(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+                self.manifest_gz)
+        m.assert_directory_verifies('')
+
+
+class CompressedSubManifestTest(TempDirTestCase):
+    """
+    Test a tree with top-level Manifest being compressed.
+    """
+
+    # we can't compress locally here since we need stable result
+    SUB_MANIFEST = b'''
+H4sICHX68FkCA01hbmlmZXN0AHNxDHFUKEktLlEwUPB1MVVIMTFMsUhOsbRIMzBIMjIwSbW0MDCw
+tLRITU6zMDEyT+UCAJqyznMxAAAA
+'''
+    DIRS = ['sub']
+    FILES = {
+        'Manifest': u'''
+MANIFEST sub/Manifest.gz 78 MD5 9c158f87b2445279d7c8aac439612fba
+''',
+        'sub/test': u'',
+    }
+
+    def setUp(self):
+        super(CompressedSubManifestTest, self).setUp()
+        self.manifest_gz = os.path.join(self.dir, 'sub/Manifest.gz')
+        with io.open(self.manifest_gz, 'wb') as f:
+            f.write(base64.b64decode(self.SUB_MANIFEST))
+
+    def tearDown(self):
+        os.unlink(self.manifest_gz)
+        super(CompressedSubManifestTest, self).tearDown()
+
+    def test_find_path_entry(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+                os.path.join(self.dir, 'Manifest'))
+        self.assertEqual(m.find_path_entry('sub/test').path, 'test')
+
+    def test_assert_directory_verifies(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+                os.path.join(self.dir, 'Manifest'))
+        m.assert_directory_verifies('')