Add support for .gz, .bz2 and .xz resources

geneve/utils/resource.py

@@ -79,3 +79,41 @@ def resource(uri, basedir=None, cachedir=None, cachefile=None, validate=None):
                     tmpdir = new_tmpdir
 
         yield tmpdir
+
+
+def _gz_compress(base_name, base_dir, **kwargs):
+    import gzip
+
+    if Path(base_dir).is_dir():
+        raise ValueError("cannot compress dirs with gzip")
+
+    with open(base_dir, "rb") as f_in:
+        with gzip.open(base_name + ".gz", "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+
+def _bz2_compress(base_name, base_dir, **kwargs):
+    import bz2
+
+    if Path(base_dir).is_dir():
+        raise ValueError("cannot compress dirs with bzip2")
+
+    with open(base_dir, "rb") as f_in:
+        with bz2.open(base_name + ".bz2", "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+
+def _xz_compress(base_name, base_dir, **kwargs):
+    import lzma
+
+    if Path(base_dir).is_dir():
+        raise ValueError("cannot compress dirs with lzma")
+
+    with open(base_dir, "rb") as f_in:
+        with lzma.open(base_name + ".xz", "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+
+shutil.register_archive_format("gz", _gz_compress, description="compress a file with gzip")
+shutil.register_archive_format("bz2", _bz2_compress, description="compress a file with bzip2")
+shutil.register_archive_format("xz", _xz_compress, description="compress a file with xz")

tests/data/test-resource.ndjson

@@ -0,0 +1,10 @@
+{"id": 0}
+{"id": 1}
+{"id": 2}
+{"id": 3}
+{"id": 4}
+{"id": 5}
+{"id": 6}
+{"id": 7}
+{"id": 8}
+{"id": 9}

tests/test_utils.py

@@ -72,19 +72,30 @@ class TestResource(unittest.TestCase):
     resource_xztar = data_dir / (resource.name + ".tar.xz")
     resource_zip = data_dir / (resource.name + ".zip")
 
+    resource2 = data_dir / "test-resource.ndjson"
+    resource2_bz = data_dir / (resource2.name + ".bz2")
+    resource2_gz = data_dir / (resource2.name + ".gz")
+    resource2_xz = data_dir / (resource2.name + ".xz")
+
     @classmethod
     def setUpClass(cls):
         make_archive(cls.resource, "bztar", root_dir=data_dir, base_dir=cls.resource.name)
         make_archive(cls.resource, "gztar", root_dir=data_dir, base_dir=cls.resource.name)
         make_archive(cls.resource, "xztar", root_dir=data_dir, base_dir=cls.resource.name)
         make_archive(cls.resource, "zip", root_dir=data_dir, base_dir=cls.resource.name)
+        make_archive(cls.resource2, "bz2", root_dir=data_dir, base_dir=cls.resource2.name)
+        make_archive(cls.resource2, "gz", root_dir=data_dir, base_dir=cls.resource2.name)
+        make_archive(cls.resource2, "xz", root_dir=data_dir, base_dir=cls.resource2.name)
 
     @classmethod
     def tearDownClass(cls):
         cls.resource_bztar.unlink()
         cls.resource_gztar.unlink()
         cls.resource_xztar.unlink()
         cls.resource_zip.unlink()
+        cls.resource2_bz.unlink()
+        cls.resource2_gz.unlink()
+        cls.resource2_xz.unlink()
 
     def test_dir(self):
         uri = str(self.resource)
@@ -95,12 +106,17 @@ def test_dir(self):
             self.assertTrue(manifest.exists(), msg=f"{manifest} does not exist")
 
     def test_local(self):
-        for ext in ["tar.bz2", "tar.gz", "tar.xz", "zip"]:
+        for ext, name in [
+            ("tar.bz2", self.resource.name),
+            ("tar.gz", self.resource.name),
+            ("tar.xz", self.resource.name),
+            ("zip", self.resource.name),
+        ]:
             tests = [
-                (f"file://./tests/data/{self.resource.name}.{ext}", None),
-                (f"file://./{self.resource.name}.{ext}", data_dir),
-                (f"tests/data/{self.resource.name}.{ext}", None),
-                (f"{self.resource.name}.{ext}", data_dir),
+                (f"file://./tests/data/{name}.{ext}", None),
+                (f"file://./{name}.{ext}", data_dir),
+                (f"tests/data/{name}.{ext}", None),
+                (f"{name}.{ext}", data_dir),
             ]
 
             for uri, basedir in tests: