add zstd support

game_backuper/zstd.py

@@ -0,0 +1,177 @@
+try:
+    from game_backuper._zstd import (
+        ZSTDCompressor,
+        ZSTDDecompressor,
+        maxCLevel,
+    )
+    have_zstd = True
+except ImportError:
+    have_zstd = False
+if have_zstd:
+    from builtins import open as _builtin_open
+    from _compression import BaseStream, DecompressReader
+    import os
+    import io
+
+    _MODE_CLOSED = 0
+    _MODE_READ = 1
+    _MODE_WRITE = 3
+    MAX_COMPRESS_LEVEL = maxCLevel()
+
+    class ZSTDFile(BaseStream):
+        def __init__(self, filename, mode="r", *, compresslevel=3):
+            self._fp = None
+            self._closefp = False
+            self._mode = _MODE_CLOSED
+            if not (1 <= compresslevel <= MAX_COMPRESS_LEVEL):
+                raise ValueError(f"compresslevel must be between 1 and {MAX_COMPRESS_LEVEL}")  # noqa: E501
+            if mode in ("", "r", "rb"):
+                mode = "rb"
+                mode_code = _MODE_READ
+            elif mode in ("w", "wb"):
+                mode = "wb"
+                mode_code = _MODE_WRITE
+                self._compressor = ZSTDCompressor(compresslevel)
+            elif mode in ("x", "xb"):
+                mode = "xb"
+                mode_code = _MODE_WRITE
+                self._compressor = ZSTDCompressor(compresslevel)
+            elif mode in ("a", "ab"):
+                mode = "ab"
+                mode_code = _MODE_WRITE
+                self._compressor = ZSTDCompressor(compresslevel)
+            else:
+                raise ValueError("Invalid mode: %r" % (mode,))
+            if isinstance(filename, (str, bytes, os.PathLike)):
+                self._fp = _builtin_open(filename, mode)
+                self._closefp = True
+                self._mode = mode_code
+            else:
+                raise TypeError("filename must be a str, bytes, file or PathLike object")  # noqa: E501
+            if self._mode == _MODE_READ:
+                raw = DecompressReader(self._fp, ZSTDDecompressor,
+                                       trailing_error=OSError)
+                self._buffer = io.BufferedReader(raw)
+            else:
+                self._pos = 0
+
+        def close(self):
+            if self._mode == _MODE_CLOSED:
+                return
+            try:
+                if self._mode == _MODE_READ:
+                    self._buffer.close()
+                elif self._mode == _MODE_WRITE:
+                    self._fp.write(self._compressor.flush())
+                    self._compressor = None
+            finally:
+                try:
+                    if self._closefp:
+                        self._fp.close()
+                finally:
+                    self._fp = None
+                    self._closefp = False
+                    self._mode = _MODE_CLOSED
+                    self._buffer = None
+
+        @property
+        def closed(self):
+            return self._mode == _MODE_CLOSED
+
+        def fileno(self):
+            self._check_not_closed()
+            return self._fp.fileno()
+
+        def seekable(self):
+            return self.readable() and self._buffer.seekable()
+
+        def readable(self):
+            self._check_not_closed()
+            return self._mode == _MODE_READ
+
+        def writable(self):
+            self._check_not_closed()
+            return self._mode == _MODE_WRITE
+
+        def peek(self, n=0):
+            self._check_can_read()
+            return self._buffer.peek(n)
+
+        def read(self, size=-1):
+            self._check_can_read()
+            return self._buffer.read(size)
+
+        def read1(self, size=-1):
+            self._check_can_read()
+            if size < 0:
+                size = io.DEFAULT_BUFFER_SIZE
+            return self._buffer.read1(size)
+
+        def readinto(self, b):
+            self._check_can_read()
+            return self._buffer.readinto(b)
+
+        def readline(self, size=-1):
+            if not isinstance(size, int):
+                if not hasattr(size, "__index__"):
+                    raise TypeError("Integer argument expected")
+                size = size.__index__()
+            self._check_can_read()
+            return self._buffer.readline(size)
+
+        def __iter__(self):
+            self._check_can_read()
+            return self._buffer.__iter__()
+
+        def readlines(self, size=-1):
+            if not isinstance(size, int):
+                if not hasattr(size, "__index__"):
+                    raise TypeError("Integer argument expected")
+                size = size.__index__()
+            self._check_can_read()
+            return self._buffer.readlines(size)
+
+        def write(self, data):
+            self._check_can_write()
+            if isinstance(data, (bytes, bytearray)):
+                length = len(data)
+            else:
+                data = memoryview(data)
+                length = data.nbytes
+            compressed = self._compressor.compress(data)
+            self._fp.write(compressed)
+            self._pos += length
+            return length
+
+        def writelines(self, seq):
+            return BaseStream.writelines(self, seq)
+
+        def seek(self, offset, whence=io.SEEK_SET):
+            self._check_can_seek()
+            return self._buffer.seek(offset, whence)
+
+        def tell(self):
+            self._check_not_closed()
+            if self._mode == _MODE_READ:
+                return self._buffer.tell()
+            return self._pos
+
+    def open(filename, mode="rb", compresslevel=3, encoding=None, errors=None, newline=None):  # noqa: E501
+        if "t" in mode:
+            if "b" in mode:
+                raise ValueError("Invalid mode: %r" % (mode,))
+        else:
+            if encoding is not None:
+                raise ValueError("Argument 'encoding' not supported in binary mode")  # noqa: E501
+            if errors is not None:
+                raise ValueError("Argument 'errors' not supported in binary mode")  # noqa: E501
+            if newline is not None:
+                raise ValueError("Argument 'newline' not supported in binary mode")  # noqa: E501
+
+        bz_mode = mode.replace("t", "")
+        binary_file = ZSTDFile(filename, bz_mode, compresslevel=compresslevel)
+
+        if "t" in mode:
+            return io.TextIOWrapper(binary_file, encoding, errors, newline)
+        else:
+            return binary_file