test streaming gz, lz4, bz2, xz and zst

Author: lhoestq

src/datasets/config.py

@@ -120,21 +120,10 @@
     logger.info("Disabling Apache Beam because USE_BEAM is set to False")
 
 
-USE_RAR = os.environ.get("USE_RAR", "AUTO").upper()
-RARFILE_VERSION = "N/A"
-RARFILE_AVAILABLE = False
-if USE_RAR in ("1", "ON", "YES", "AUTO"):
-    try:
-        RARFILE_VERSION = version.parse(importlib_metadata.version("rarfile"))
-        RARFILE_AVAILABLE = True
-        logger.info("rarfile available.")
-    except importlib_metadata.PackageNotFoundError:
-        pass
-else:
-    logger.info("Disabling rarfile because USE_RAR is set to False")
-
-
+# Optional compression tools
+RARFILE_AVAILABLE = importlib.util.find_spec("rarfile") is not None
 ZSTANDARD_AVAILABLE = importlib.util.find_spec("zstandard") is not None
+LZ4_AVAILABLE = importlib.util.find_spec("lz4") is not None
 
 
 # Cache location

src/datasets/filesystems/compression.py

@@ -39,7 +39,7 @@ def __init__(
         )
         self.info = self.file.fs.info(self.file.path)
         self.compressed_name = os.path.basename(self.file.path.split("::")[0])
-        self.uncompressed_name = self.compressed_name.rstrip(self.extension)
+        self.uncompressed_name = self.compressed_name[: self.compressed_name.rindex(".")]
         self.dir_cache = None
 
     @classmethod

tests/conftest.py

@@ -7,6 +7,7 @@
 import pyarrow.parquet as pq
 import pytest
 
+from datasets import config
 from datasets.arrow_dataset import Dataset
 from datasets.features import ClassLabel, Features, Sequence, Value
 
@@ -80,24 +81,59 @@ def text_file(tmp_path_factory):
 
 @pytest.fixture(scope="session")
 def xz_file(tmp_path_factory):
-    filename = tmp_path_factory.mktemp("data") / "file.xz"
+    filename = tmp_path_factory.mktemp("data") / "file.txt.xz"
     data = bytes(FILE_CONTENT, "utf-8")
     with lzma.open(filename, "wb") as f:
         f.write(data)
     return filename
 
 
 @pytest.fixture(scope="session")
-def gz_path(tmp_path_factory, text_path):
+def gz_file(tmp_path_factory):
     import gzip
 
-    path = str(tmp_path_factory.mktemp("data") / "file.gz")
+    path = str(tmp_path_factory.mktemp("data") / "file.txt.gz")
     data = bytes(FILE_CONTENT, "utf-8")
     with gzip.open(path, "wb") as f:
         f.write(data)
     return path
 
 
+@pytest.fixture(scope="session")
+def bz2_file(tmp_path_factory):
+    import bz2
+
+    path = tmp_path_factory.mktemp("data") / "file.txt.bz2"
+    data = bytes(FILE_CONTENT, "utf-8")
+    with bz2.open(path, "wb") as f:
+        f.write(data)
+    return path
+
+
+@pytest.fixture(scope="session")
+def zstd_file(tmp_path_factory):
+    if config.ZSTANDARD_AVAILABLE:
+        import zstandard as zstd
+
+        path = tmp_path_factory.mktemp("data") / "file.txt.zst"
+        data = bytes(FILE_CONTENT, "utf-8")
+        with zstd.open(path, "wb") as f:
+            f.write(data)
+        return path
+
+
+@pytest.fixture(scope="session")
+def lz4_file(tmp_path_factory):
+    if config.LZ4_AVAILABLE:
+        import lz4.frame
+
+        path = tmp_path_factory.mktemp("data") / "file.txt.lz4"
+        data = bytes(FILE_CONTENT, "utf-8")
+        with lz4.frame.open(path, "wb") as f:
+            f.write(data)
+        return path
+
+
 @pytest.fixture(scope="session")
 def xml_file(tmp_path_factory):
     filename = tmp_path_factory.mktemp("data") / "file.xml"

tests/test_extract.py

@@ -1,25 +1,13 @@
 import pytest
-import zstandard as zstd
 
 from datasets.utils.extract import Extractor, ZstdExtractor
 
+from .utils import require_zstandard
 
-FILE_CONTENT = """\
-    Text data.
-    Second line of data."""
 
-
-@pytest.fixture(scope="session")
-def zstd_path(tmp_path_factory):
-    path = tmp_path_factory.mktemp("data") / "file.zstd"
-    data = bytes(FILE_CONTENT, "utf-8")
-    with zstd.open(path, "wb") as f:
-        f.write(data)
-    return path
-
-
-def test_zstd_extractor(zstd_path, tmp_path, text_file):
-    input_path = zstd_path
+@require_zstandard
+def test_zstd_extractor(zstd_file, tmp_path, text_file):
+    input_path = zstd_file
     assert ZstdExtractor.is_extractable(input_path)
     output_path = str(tmp_path / "extracted.txt")
     ZstdExtractor.extract(input_path, output_path)
@@ -30,21 +18,16 @@ def test_zstd_extractor(zstd_path, tmp_path, text_file):
     assert extracted_file_content == expected_file_content
 
 
-@pytest.mark.parametrize(
-    "compression_format, expected_text_path_name", [("gzip", "text_path"), ("xz", "text_file"), ("zstd", "text_file")]
-)
-def test_extractor(
-    compression_format, expected_text_path_name, text_gz_path, xz_file, zstd_path, tmp_path, text_file, text_path
-):
-    input_paths = {"gzip": text_gz_path, "xz": xz_file, "zstd": zstd_path}
+@require_zstandard
+@pytest.mark.parametrize("compression_format", ["gzip", "xz", "zstd"])
+def test_extractor(compression_format, gz_file, xz_file, zstd_file, tmp_path, text_file):
+    input_paths = {"gzip": gz_file, "xz": xz_file, "zstd": zstd_file}
     input_path = str(input_paths[compression_format])
     output_path = str(tmp_path / "extracted.txt")
     assert Extractor.is_extractable(input_path)
     Extractor.extract(input_path, output_path)
     with open(output_path) as f:
         extracted_file_content = f.read()
-    expected_text_paths = {"text_file": text_file, "text_path": text_path}
-    expected_text_path = str(expected_text_paths[expected_text_path_name])
-    with open(expected_text_path) as f:
+    with open(text_file) as f:
         expected_file_content = f.read()
     assert extracted_file_content == expected_file_content

tests/test_streaming_download_manager.py

@@ -2,7 +2,9 @@
 
 import pytest
 
-from .utils import require_streaming
+from datasets.filesystems import COMPRESSION_FILESYSTEMS
+
+from .utils import require_lz4, require_streaming, require_zstandard
 
 
 TEST_URL = "https://huggingface.co/datasets/lhoestq/test/raw/main/some_text.txt"
@@ -116,3 +118,24 @@ def test_streaming_dl_manager_download_and_extract_with_join(input_path, filenam
     extracted_path = dl_manager.download_and_extract(input_path)
     output_path = xjoin(extracted_path, filename)
     assert output_path == expected_path
+
+
+@require_streaming
+@require_zstandard
+@require_lz4
+@pytest.mark.parametrize("compression_fs_class", COMPRESSION_FILESYSTEMS)
+def test_streaming_dl_manager_extract_all_supported_single_file_compression_types(
+    compression_fs_class, gz_file, xz_file, zstd_file, bz2_file, lz4_file, text_file
+):
+    from datasets.utils.streaming_download_manager import StreamingDownloadManager, xopen
+
+    input_paths = {"gzip": gz_file, "xz": xz_file, "zstd": zstd_file, "bz2": bz2_file, "lz4": lz4_file}
+    input_path = str(input_paths[compression_fs_class.protocol])
+    dl_manager = StreamingDownloadManager()
+    output_path = dl_manager.extract(input_path)
+    path = os.path.basename(input_path)
+    path = path[: path.rindex(".")]
+    assert output_path == f"{compression_fs_class.protocol}://{path}::{input_path}"
+    fsspec_open_file = xopen(output_path, encoding="utf-8")
+    with fsspec_open_file as f, open(text_file, encoding="utf-8") as expected_file:
+        assert f.read() == expected_file.read()

tests/utils.py

@@ -137,6 +137,30 @@ def require_jax(test_case):
     return test_case
 
 
+def require_zstandard(test_case):
+    """
+    Decorator marking a test that requires zstandard.
+
+    These tests are skipped when zstandard isn't installed.
+
+    """
+    if not config.ZSTANDARD_AVAILABLE:
+        test_case = unittest.skip("test requires zstandard")(test_case)
+    return test_case
+
+
+def require_lz4(test_case):
+    """
+    Decorator marking a test that requires lz4.
+
+    These tests are skipped when lz4 isn't installed.
+
+    """
+    if not config.LZ4_AVAILABLE:
+        test_case = unittest.skip("test requires lz4")(test_case)
+    return test_case
+
+
 def require_transformers(test_case):
     """
     Decorator marking a test that requires transformers.