510 lines
16 KiB
Python
510 lines
16 KiB
Python
try:
|
|
from bz2 import BZ2File, BZ2Compressor, BZ2Decompressor
|
|
have_bz2 = True
|
|
except ImportError:
|
|
have_bz2 = False
|
|
try:
|
|
from gzip import GzipFile
|
|
have_gzip = True
|
|
except ImportError:
|
|
have_gzip = False
|
|
try:
|
|
from lzma import LZMAFile, LZMACompressor, LZMADecompressor
|
|
have_lzma = True
|
|
except ImportError:
|
|
have_lzma = False
|
|
try:
|
|
from lzip import (
|
|
FileEncoder as LZIPFileEncoder,
|
|
decompress_file_iter as LZIP_decompress_file_iter,
|
|
level_to_dictionary_size_and_match_len_limit as lzip_level_dict,
|
|
)
|
|
from lzip_extension import (
|
|
Decoder as LZIPDecoder,
|
|
Encoder as LZIPEncoder,
|
|
)
|
|
have_lzip = True
|
|
except ImportError:
|
|
have_lzip = False
|
|
try:
|
|
from game_backuper.zstd import (
|
|
ZSTDCompressor,
|
|
ZSTDDecompressor,
|
|
ZSTDFile,
|
|
MAX_COMPRESS_LEVEL as ZSTD_MAX,
|
|
)
|
|
have_zstd = True
|
|
except ImportError:
|
|
have_zstd = False
|
|
try:
|
|
from snappy import (
|
|
StreamCompressor as Snappy_Compressor,
|
|
StreamDecompressor as Snappy_Decompressor,
|
|
)
|
|
have_snappy = True
|
|
except ImportError:
|
|
have_snappy = False
|
|
try:
|
|
from brotli import (
|
|
Compressor as BrotliCompressor,
|
|
Decompressor as BrotliDecompressor,
|
|
)
|
|
have_brotli = True
|
|
except ImportError:
|
|
have_brotli = False
|
|
from enum import IntEnum, unique
|
|
try:
|
|
from functools import cached_property
|
|
except ImportError:
|
|
cached_property = property
|
|
from os.path import exists, isfile, getsize
|
|
from os import remove
|
|
from game_backuper.file import mkdir_for_file, hydrate_file_if_needed
|
|
|
|
|
|
if have_gzip:
|
|
class GZCompressor:
|
|
def __init__(self, compress_level: int = 9, file_obj=None):
|
|
self.write_to_file = True
|
|
self._fileobj = file_obj
|
|
self._level = compress_level
|
|
self._f = None
|
|
|
|
def close(self):
|
|
if self._f is None:
|
|
return
|
|
self._f.close()
|
|
|
|
def compress(self, data: bytes) -> bytes:
|
|
if self._f is None:
|
|
self._f = GzipFile(mode="wb", compresslevel=self._level, fileobj=self._fileobj) # noqa: E501
|
|
self._f.write(data)
|
|
return b''
|
|
|
|
def flush(self) -> bytes:
|
|
self._f.flush()
|
|
return b''
|
|
|
|
class GZDecompressor:
|
|
def __init__(self, file_obj=None):
|
|
self.write_to_file = True
|
|
self._fileobj = file_obj
|
|
self._f = None
|
|
|
|
def close(self):
|
|
if self._f is None:
|
|
return
|
|
self._f.close()
|
|
|
|
def read(self, len: int) -> bytes:
|
|
if self._f is None:
|
|
self._f = GzipFile(mode="rb", fileobj=self._fileobj)
|
|
return self._f.read(len)
|
|
|
|
if have_lzip:
|
|
class LZIPCompressor:
|
|
def __init__(self, level: int = 6):
|
|
d = lzip_level_dict[level]
|
|
self._encoder = LZIPEncoder(d[0], d[1], 1 << 51)
|
|
|
|
def compress(self, data: bytes) -> bytes:
|
|
return self._encoder.compress(data)
|
|
|
|
def flush(self) -> bytes:
|
|
return self._encoder.finish()
|
|
|
|
class LZIPDecompressor:
|
|
def __init__(self):
|
|
self._decoder = LZIPDecoder(1)
|
|
|
|
def decompress(self, data: bytes) -> bytes:
|
|
return self._decoder.decompress(data)
|
|
|
|
@property
|
|
def eof(self):
|
|
return False
|
|
|
|
if have_brotli:
|
|
class BrotliCompressor2:
|
|
def __init__(self, *k, **kw):
|
|
self._c = BrotliCompressor(*k, **kw)
|
|
|
|
def compress(self, data: bytes) -> bytes:
|
|
return self._c.process(data)
|
|
|
|
def flush(self) -> bytes:
|
|
return self._c.finish()
|
|
|
|
class BrotliDecompressor2:
|
|
def __init__(self, *k, **kw) -> None:
|
|
self._c = BrotliDecompressor(*k, **kw)
|
|
|
|
def decompress(self, data: bytes) -> bytes:
|
|
return self._c.process(data)
|
|
|
|
@property
|
|
def eof(self):
|
|
return self._c.is_finished()
|
|
|
|
|
|
@unique
|
|
class CompressMethod(IntEnum):
|
|
BZIP2 = 0
|
|
GZIP = 1
|
|
LZMA = 2
|
|
LZIP = 3
|
|
ZSTD = 4
|
|
SNAPPY = 5
|
|
BROTLI = 6
|
|
|
|
@staticmethod
|
|
def from_str(v: str) -> IntEnum:
|
|
if isinstance(v, str):
|
|
t = v.lower()
|
|
if t == 'bzip2':
|
|
return CompressMethod.BZIP2
|
|
elif t == "gzip":
|
|
return CompressMethod.GZIP
|
|
elif t == "lzma":
|
|
return CompressMethod.LZMA
|
|
elif t == "lzip":
|
|
return CompressMethod.LZIP
|
|
elif t == "zstd":
|
|
return CompressMethod.ZSTD
|
|
elif t == "snappy":
|
|
return CompressMethod.SNAPPY
|
|
elif t == "brotli":
|
|
return CompressMethod.BROTLI
|
|
else:
|
|
raise TypeError('Must be str.')
|
|
|
|
def to_str(self) -> str:
|
|
return {
|
|
CompressMethod.BZIP2: 'bzip2',
|
|
CompressMethod.GZIP: 'gzip',
|
|
CompressMethod.LZMA: 'lzma',
|
|
CompressMethod.LZIP: 'lzip',
|
|
CompressMethod.ZSTD: 'zstd',
|
|
CompressMethod.SNAPPY: 'snappy',
|
|
CompressMethod.BROTLI: 'brotli',
|
|
}[self]
|
|
|
|
|
|
class CompressConfig:
|
|
def __init__(self, method: str, level: int = None):
|
|
self._method = CompressMethod.from_str(method)
|
|
if self._method is None:
|
|
raise ValueError('Unknown compress method.')
|
|
if self._method == CompressMethod.BZIP2:
|
|
if not have_bz2:
|
|
raise NotImplementedError("bzip2 not supported.")
|
|
if level is None:
|
|
self._level = 9
|
|
else:
|
|
if isinstance(level, int) and level >= 1 and level <= 9:
|
|
self._level = level
|
|
else:
|
|
raise ValueError('bzip2: compress_level should be 1-9.')
|
|
self._ext = ".bz2"
|
|
elif self._method == CompressMethod.GZIP:
|
|
if not have_gzip:
|
|
raise NotImplementedError("gzip not supported.")
|
|
if level is None:
|
|
self._level = 9
|
|
else:
|
|
if isinstance(level, int) and level >= 0 and level <= 9:
|
|
self._level = level
|
|
else:
|
|
raise ValueError('gzip: compress_level should be 0-9.')
|
|
self._ext = '.gz'
|
|
elif self._method == CompressMethod.LZMA:
|
|
if not have_lzma:
|
|
raise NotImplementedError("lzma not supported.")
|
|
if level is None:
|
|
self._level = 6
|
|
else:
|
|
if isinstance(level, int) and level >= 0 and level <= 9:
|
|
self._level = level
|
|
else:
|
|
raise ValueError('lzma: compress_level should be 0-9.')
|
|
self._ext = ".xz"
|
|
elif self._method == CompressMethod.LZIP:
|
|
if not have_lzip:
|
|
raise NotImplementedError("lzip not supported.")
|
|
if level is None:
|
|
self._level = 6
|
|
else:
|
|
if isinstance(level, int) and level >= 0 and level <= 9:
|
|
self._level = level
|
|
else:
|
|
raise ValueError('lzip: compress_level should be 0-9.')
|
|
self._ext = ".lz"
|
|
elif self._method == CompressMethod.ZSTD:
|
|
if not have_zstd:
|
|
raise NotImplementedError("zstd not supported.")
|
|
if level is None:
|
|
self._level = 3
|
|
else:
|
|
if isinstance(level, int) and level >= 0 and level <= ZSTD_MAX:
|
|
self._level = level
|
|
else:
|
|
raise ValueError(f'zstd: compress_level should be 0-{ZSTD_MAX}.') # noqa: E501
|
|
self._ext = ".zst"
|
|
elif self._method == CompressMethod.SNAPPY:
|
|
if not have_snappy:
|
|
raise NotImplementedError("snappy not supported.")
|
|
self._level = None
|
|
self._ext = ".snappy"
|
|
elif self._method == CompressMethod.BROTLI:
|
|
if not have_brotli:
|
|
raise NotImplementedError("brotli not supported.")
|
|
if level is None:
|
|
self._level = None
|
|
else:
|
|
if isinstance(level, int) and level >= 0 and level <= 11:
|
|
self._level = level
|
|
else:
|
|
raise ValueError('brotli: compress_level should be 0-11.')
|
|
self._ext = '.br'
|
|
self._chunk_size = 131072
|
|
|
|
def __repr__(self):
|
|
t = type(self)
|
|
return f"<{t.__module__}.{t.__qualname__} object at {hex(id(self))}; method={repr(self._method)}, level={repr(self._level)}, ext={repr(self._ext)}>" # noqa: E501
|
|
|
|
def compressor(self, fileobj):
|
|
if self._method == CompressMethod.BZIP2:
|
|
return BZ2Compressor(self._level)
|
|
elif self._method == CompressMethod.GZIP:
|
|
return GZCompressor(self._level, fileobj)
|
|
elif self._method == CompressMethod.LZMA:
|
|
return LZMACompressor(preset=self._level)
|
|
elif self._method == CompressMethod.LZIP:
|
|
return LZIPCompressor(self._level)
|
|
elif self._method == CompressMethod.ZSTD:
|
|
return ZSTDCompressor(self._level)
|
|
elif self._method == CompressMethod.SNAPPY:
|
|
return Snappy_Compressor()
|
|
elif self._method == CompressMethod.BROTLI:
|
|
c = {}
|
|
if self._level is not None:
|
|
c['quality'] = self._level
|
|
return BrotliCompressor2(**c)
|
|
|
|
def decompressor(self, fileobj):
|
|
if self._method == CompressMethod.BZIP2:
|
|
return BZ2Decompressor()
|
|
elif self._method == CompressMethod.GZIP:
|
|
return GZDecompressor(fileobj)
|
|
elif self._method == CompressMethod.LZMA:
|
|
return LZMADecompressor()
|
|
elif self._method == CompressMethod.LZIP:
|
|
return LZIPDecompressor()
|
|
elif self._method == CompressMethod.ZSTD:
|
|
return ZSTDDecompressor()
|
|
elif self._method == CompressMethod.SNAPPY:
|
|
return Snappy_Decompressor()
|
|
elif self._method == CompressMethod.BROTLI:
|
|
return BrotliDecompressor2()
|
|
|
|
@cached_property
|
|
def chunk_size(self) -> int:
|
|
return self._chunk_size
|
|
|
|
@cached_property
|
|
def ext(self) -> str:
|
|
return self._ext
|
|
|
|
@cached_property
|
|
def level(self) -> int:
|
|
return self._level
|
|
|
|
@cached_property
|
|
def method(self) -> CompressMethod:
|
|
return self._method
|
|
|
|
|
|
supported_exts = []
|
|
if have_bz2:
|
|
supported_exts.append('.bz2')
|
|
if have_gzip:
|
|
supported_exts.append('.gz')
|
|
if have_lzma:
|
|
supported_exts.append('.xz')
|
|
if have_lzip:
|
|
supported_exts.append('.lz')
|
|
if have_zstd:
|
|
supported_exts.append('.zst')
|
|
if have_snappy:
|
|
supported_exts.append('.snappy')
|
|
if have_brotli:
|
|
supported_exts.append('.br')
|
|
|
|
|
|
def sizeof_fmt(num, suffix='B'):
|
|
for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', ' Ei', 'Zi']:
|
|
if abs(num) < 1024.0:
|
|
return "%3.1f%s%s" % (num, unit, suffix)
|
|
num /= 1024.0
|
|
return "%.1f%s%s" % (num, 'Yi', suffix)
|
|
|
|
|
|
def compress_info(ori: int, re: int):
|
|
if ori == 0:
|
|
return f"{sizeof_fmt(ori)} -> {sizeof_fmt(re)} ({float('inf')})"
|
|
return f"{sizeof_fmt(ori)} -> {sizeof_fmt(re)} ({re/ori*100:.2f}%)"
|
|
|
|
|
|
def compress(src: str, dest: str, c: CompressConfig, name: str, prog: str):
|
|
exts = [''] + supported_exts.copy()
|
|
exts.remove(c.ext)
|
|
fn = dest + c.ext
|
|
cs = c.chunk_size
|
|
if exists(fn):
|
|
remove(fn)
|
|
mkdir_for_file(fn)
|
|
if c.method == CompressMethod.BZIP2:
|
|
with open(src, 'rb') as t:
|
|
with BZ2File(fn, 'wb', compresslevel=c.level) as f:
|
|
a = t.read(cs)
|
|
while a != b'':
|
|
f.write(a)
|
|
a = t.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.GZIP:
|
|
with open(src, 'rb') as t:
|
|
with GzipFile(fn, 'wb', compresslevel=c.level) as f:
|
|
a = t.read(cs)
|
|
while a != b'':
|
|
f.write(a)
|
|
a = t.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.LZMA:
|
|
with open(src, 'rb') as t:
|
|
with LZMAFile(fn, 'wb', preset=c.level) as f:
|
|
a = t.read(cs)
|
|
while a != b'':
|
|
f.write(a)
|
|
a = t.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.LZIP:
|
|
with open(src, 'rb') as t:
|
|
with LZIPFileEncoder(fn, c.level) as f:
|
|
a = t.read(cs)
|
|
while a != b'':
|
|
f.compress(a)
|
|
a = t.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.ZSTD:
|
|
with open(src, 'rb') as t:
|
|
with ZSTDFile(fn, 'wb', compresslevel=c.level) as f:
|
|
a = t.read(cs)
|
|
while a != b'':
|
|
f.write(a)
|
|
a = t.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.SNAPPY:
|
|
with open(src, 'rb') as t:
|
|
with open(fn, 'wb') as f:
|
|
o = Snappy_Compressor()
|
|
a = t.read(cs)
|
|
while a != b'':
|
|
b = o.compress(a)
|
|
f.write(b)
|
|
a = t.read(cs)
|
|
del a, b, o
|
|
elif c.method == CompressMethod.BROTLI:
|
|
k = {}
|
|
if c.level is not None:
|
|
k['quality'] = c.level
|
|
with open(src, 'rb') as t:
|
|
with open(fn, 'wb') as f:
|
|
o = BrotliCompressor(**k)
|
|
a = t.read(cs)
|
|
while a != b'':
|
|
b = o.process(a)
|
|
f.write(b)
|
|
a = t.read(cs)
|
|
f.write(o.finish())
|
|
del a, b, o
|
|
del k
|
|
i = compress_info(getsize(src), getsize(fn))
|
|
print(f'{prog}: Compressed {src}({name}) -> {fn} ({i})')
|
|
del i
|
|
for i in exts:
|
|
f = dest + i
|
|
if exists(f) and isfile(f):
|
|
remove(f)
|
|
print(f'{prog}: Removed {f}({name})')
|
|
|
|
|
|
def decompress(src: str, dest: str, c: CompressConfig, name: str, prog: str):
|
|
fn = src + c.ext
|
|
if exists(dest):
|
|
remove(dest)
|
|
cs = c.chunk_size
|
|
hydrate_file_if_needed(fn)
|
|
if c.method == CompressMethod.BZIP2:
|
|
with BZ2File(fn, 'rb') as f:
|
|
with open(dest, 'wb') as t:
|
|
a = f.read(cs)
|
|
while a != b'':
|
|
t.write(a)
|
|
a = f.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.GZIP:
|
|
with GzipFile(fn, 'rb') as f:
|
|
with open(dest, 'wb') as t:
|
|
a = f.read(cs)
|
|
while a != b'':
|
|
t.write(a)
|
|
a = f.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.LZMA:
|
|
with LZMAFile(fn, 'rb') as f:
|
|
with open(dest, 'wb') as t:
|
|
a = f.read(cs)
|
|
while a != b'':
|
|
t.write(a)
|
|
a = f.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.LZIP:
|
|
with open(dest, 'wb') as t:
|
|
f = LZIP_decompress_file_iter(fn, chunk_size=cs)
|
|
for a in f:
|
|
t.write(a)
|
|
del a
|
|
elif c.method == CompressMethod.ZSTD:
|
|
with ZSTDFile(fn, 'rb') as f:
|
|
with open(dest, 'wb') as t:
|
|
a = f.read(cs)
|
|
while a != b'':
|
|
t.write(a)
|
|
a = f.read(cs)
|
|
del a
|
|
elif c.method == CompressMethod.SNAPPY:
|
|
with open(fn, 'rb') as f:
|
|
with open(dest, 'wb') as t:
|
|
o = Snappy_Decompressor()
|
|
a = f.read(cs)
|
|
while a != b'':
|
|
b = o.decompress(a)
|
|
t.write(b)
|
|
a = f.read(cs)
|
|
o.flush()
|
|
del a, b, o
|
|
elif c.method == CompressMethod.BROTLI:
|
|
with open(fn, 'rb') as f:
|
|
with open(dest, 'wb') as t:
|
|
o = BrotliDecompressor()
|
|
a = f.read(cs)
|
|
while a != b'':
|
|
b = o.process(a)
|
|
t.write(b)
|
|
a = f.read(cs)
|
|
if not o.is_finished():
|
|
raise ValueError('Read all datas from file but seems not finished.') # noqa: E501
|
|
del a, b, o
|
|
i = compress_info(getsize(dest), getsize(fn))
|
|
print(f'{prog}: Decompressed {fn}({name}) -> {dest} ({i})')
|