Replace zstd with stdlib and backports zstd

Author: Schamper

dissect/hypervisor/disk/qcow2.py

@@ -3,9 +3,9 @@
 # - https://github.com/qemu/qemu/blob/master/docs/interop/qcow2.txt
 from __future__ import annotations
 
+import sys
 import zlib
 from functools import cached_property, lru_cache
-from io import BytesIO
 from pathlib import Path
 from typing import TYPE_CHECKING, BinaryIO
 
@@ -28,8 +28,10 @@
     from collections.abc import Iterator
 
 try:
-    import zstandard as zstd
-
+    if sys.version_info >= (3, 14):
+        from compression import zstd
+    else:
+        from backports import zstd
     HAS_ZSTD = True
 except ImportError:
     HAS_ZSTD = False
@@ -384,16 +386,8 @@ def _decompress(self, buf: bytes) -> bytes:
             return dctx.decompress(buf, self.qcow2.cluster_size)
 
         if self.qcow2.compression_type == c_qcow2.QCOW2_COMPRESSION_TYPE_ZSTD:
-            result = []
-
             dctx = zstd.ZstdDecompressor()
-            reader = dctx.stream_reader(BytesIO(buf))
-            while reader.tell() < self.qcow2.cluster_size:
-                chunk = reader.read(self.qcow2.cluster_size - reader.tell())
-                if not chunk:
-                    break
-                result.append(chunk)
-            return b"".join(result)
+            return dctx.decompress(buf, self.qcow2.cluster_size)
 
         raise Error(f"Invalid compression type: {self.qcow2.compression_type}")
 

pyproject.toml

@@ -39,6 +39,7 @@ repository = "https://github.com/fox-it/dissect.hypervisor"
 [project.optional-dependencies]
 full = [
     "pycryptodome",
+    "backports.zstd; python_version < '3.14'",
 ]
 dev = [
     "dissect.hypervisor[full]",

tests/_data/disk/qcow2/basic-zstd.qcow2.gz

@@ -89,6 +89,11 @@ def basic_qcow2() -> Iterator[BinaryIO]:
     yield from open_file_gz("_data/disk/qcow2/basic.qcow2.gz")
 
 
+@pytest.fixture
+def basic_zstd_qcow2() -> Iterator[BinaryIO]:
+    yield from open_file_gz("_data/disk/qcow2/basic-zstd.qcow2.gz")
+
+
 @pytest.fixture
 def data_file_qcow2() -> Path:
     return absolute_path("_data/disk/qcow2/data-file.qcow2.gz")

tests/conftest.py

@@ -16,8 +16,9 @@ def mock_open_gz(self: Path, *args, **kwargs) -> BinaryIO:
     return gzip.open(self if self.suffix.lower() == ".gz" else self.with_suffix(self.suffix + ".gz"))
 
 
-def test_basic(basic_qcow2: BinaryIO) -> None:
-    qcow2 = QCow2(basic_qcow2)
+@pytest.mark.parametrize("name", ["basic_qcow2", "basic_zstd_qcow2"])
+def test_basic(name: str, request: pytest.FixtureRequest) -> None:
+    qcow2 = QCow2(request.getfixturevalue(name))
 
     assert qcow2.backing_file is None
     assert qcow2.data_file is qcow2.fh