Allow specifying a compresslevel to ZstdTarFile.zstopen

1 files changed, 14 insertions, 5 deletions

diff --git a/linuxnamespaces/tarutils.py b/linuxnamespaces/tarutils.py
index 6285d5a..1bbfab4 100644
--- a/linuxnamespaces/tarutils.py
+++ b/linuxnamespaces/tarutils.py
@@ -31,6 +31,7 @@ class ZstdTarFile(tarfile.TarFile):
         name: str,
         mode: typing.Literal["r", "w", "x"] = "r",
         fileobj: typing.BinaryIO | None = None,
+        compresslevel: int | None = None,
         **kwargs: typing.Any,
     ) -> tarfile.TarFile:
         if mode not in ("r", "w", "x"):
@@ -45,11 +46,19 @@ class ZstdTarFile(tarfile.TarFile):
         if mode == "r":
             zfobj = zstandard.open(openobj, "rb")
         else:
-            zfobj = zstandard.open(
-                openobj,
-                mode + "b",
-                cctx=zstandard.ZstdCompressor(write_checksum=True, threads=-1),
-            )
+            if compresslevel is not None:
+                if compresslevel > 22:
+                    raise ValueError(
+                        "invalid compression level {compresslevel}"
+                    )
+                cctx = zstandard.ZstdCompressor(
+                    write_checksum=True, threads=-1, level=compresslevel
+                )
+            else:
+                cctx = zstandard.ZstdCompressor(
+                    write_checksum=True, threads=-1
+                )
+            zfobj = zstandard.open(openobj, mode + "b", cctx=cctx)
         try:
             tarobj = cls.taropen(name, mode, zfobj, **kwargs)
         except (OSError, EOFError, zstandard.ZstdError) as exc: