Add zstandard compression support for project export

gns3server/utils/zipfile_zstd/__init__.py

@@ -0,0 +1,10 @@
+
+# NOTE: this patches the standard zipfile module
+from . import _zipfile
+
+from zipfile import *
+from zipfile import (
+    ZIP_ZSTANDARD,
+    ZSTANDARD_VERSION,
+)
+

gns3server/utils/zipfile_zstd/_patcher.py

@@ -0,0 +1,20 @@
+import functools
+
+
+class patch:
+
+    originals = {}
+
+    def __init__(self, host, name):
+        self.host = host
+        self.name = name
+
+    def __call__(self, func):
+        original = getattr(self.host, self.name)
+        self.originals[self.name] = original
+
+        functools.update_wrapper(func, original)
+        setattr(self.host, self.name, func)
+
+        return func
+

gns3server/utils/zipfile_zstd/_zipfile.py

@@ -0,0 +1,64 @@
+import zipfile
+import zstandard as zstd
+import inspect
+
+from ._patcher import patch
+
+
+zipfile.ZIP_ZSTANDARD = 93
+zipfile.compressor_names[zipfile.ZIP_ZSTANDARD] = 'zstandard'
+zipfile.ZSTANDARD_VERSION = 20
+
+
+@patch(zipfile, '_check_compression')
+def zstd_check_compression(compression):
+    if compression == zipfile.ZIP_ZSTANDARD:
+        pass
+    else:
+        patch.originals['_check_compression'](compression)
+
+
+class ZstdDecompressObjWrapper:
+    def __init__(self, o):
+        self.o = o
+
+    def __getattr__(self, attr):
+        if attr == 'eof':
+            return False
+        return getattr(self.o, attr)
+
+
+@patch(zipfile, '_get_decompressor')
+def zstd_get_decompressor(compress_type):
+    if compress_type == zipfile.ZIP_ZSTANDARD:
+        return ZstdDecompressObjWrapper(zstd.ZstdDecompressor(max_window_size=2147483648).decompressobj())
+    else:
+        return patch.originals['_get_decompressor'](compress_type)
+
+
+if 'compresslevel' in inspect.signature(zipfile._get_compressor).parameters:
+    @patch(zipfile, '_get_compressor')
+    def zstd_get_compressor(compress_type, compresslevel=None):
+        if compress_type == zipfile.ZIP_ZSTANDARD:
+            if compresslevel is None:
+                compresslevel = 3
+            return zstd.ZstdCompressor(level=compresslevel, threads=12).compressobj()
+        else:
+            return patch.originals['_get_compressor'](compress_type, compresslevel=compresslevel)
+else:
+    @patch(zipfile, '_get_compressor')
+    def zstd_get_compressor(compress_type, compresslevel=None):
+        if compress_type == zipfile.ZIP_ZSTANDARD:
+            if compresslevel is None:
+                compresslevel = 3
+            return zstd.ZstdCompressor(level=compresslevel, threads=12).compressobj()
+        else:
+            return patch.originals['_get_compressor'](compress_type)
+
+
+@patch(zipfile.ZipInfo, 'FileHeader')
+def zstd_FileHeader(self, zip64=None):
+    if self.compress_type == zipfile.ZIP_ZSTANDARD:
+        self.create_version = max(self.create_version, zipfile.ZSTANDARD_VERSION)
+        self.extract_version = max(self.extract_version, zipfile.ZSTANDARD_VERSION)
+    return patch.originals['FileHeader'](self, zip64=zip64)