From 4815acb4258e984011fd8dbf4592cabcdc3a3c7f Mon Sep 17 00:00:00 2001 From: lifegpc Date: Sun, 12 Sep 2021 21:12:46 +0800 Subject: [PATCH] add _zstd cython module --- game_backuper/_Python.pxd | 4 +- game_backuper/_pcre2.pyx | 4 - game_backuper/_zstd.pxd | 46 +++++++++ game_backuper/_zstd.pyx | 194 ++++++++++++++++++++++++++++++++++++++ game_backuper/compress.py | 2 +- setup.py | 10 +- 6 files changed, 252 insertions(+), 8 deletions(-) create mode 100644 game_backuper/_zstd.pxd create mode 100644 game_backuper/_zstd.pyx diff --git a/game_backuper/_Python.pxd b/game_backuper/_Python.pxd index 27739a0..271ff2d 100644 --- a/game_backuper/_Python.pxd +++ b/game_backuper/_Python.pxd @@ -1,5 +1,5 @@ cdef extern from "Python.h": - void Py_INCREF(object o) - void Py_DECREF(object o) const char* PyUnicode_AsUTF8(object unicode) const char* PyUnicode_AsUTF8AndSize(object unicode, Py_ssize_t *size) + int PyBytes_AsStringAndSize(object obj, char **buff, Py_ssize_t *length) + object PyBytes_FromStringAndSize(const char* v, Py_ssize_t le) diff --git a/game_backuper/_pcre2.pyx b/game_backuper/_pcre2.pyx index e4dc556..8336151 100644 --- a/game_backuper/_pcre2.pyx +++ b/game_backuper/_pcre2.pyx @@ -113,8 +113,6 @@ cdef class Match: if self.data != NULL: pcre2_match_data_free(self.data) self.data = NULL - Py_DECREF(self.r) - Py_DECREF(self.inp) def __getitem__(self, uint32_t i): if self.data == NULL: @@ -149,9 +147,7 @@ cdef class Match: def __init__(self, unicode inp, r): self.inp = inp - Py_INCREF(inp) self.r = r - Py_INCREF(r) def end(self) -> int: if self.data == NULL: diff --git a/game_backuper/_zstd.pxd b/game_backuper/_zstd.pxd new file mode 100644 index 0000000..ba9def3 --- /dev/null +++ b/game_backuper/_zstd.pxd @@ -0,0 +1,46 @@ +from libc.stddef cimport size_t + + +cdef extern from "zstd.h": + ctypedef struct ZSTD_CCtx: + pass + ctypedef ZSTD_CCtx ZSTD_CStream + ctypedef enum ZSTD_EndDirective: + ZSTD_e_continue = 0 + ZSTD_e_flush = 1 + ZSTD_e_end = 2 + cdef struct ZSTD_inBuffer_s: + const void* src + size_t size + size_t pos + ctypedef ZSTD_inBuffer_s ZSTD_inBuffer + cdef struct ZSTD_outBuffer_s: + void* dst + size_t size + size_t pos + ctypedef ZSTD_outBuffer_s ZSTD_outBuffer + ctypedef enum ZSTD_cParameter: + ZSTD_c_compressionLevel + ZSTD_c_checksumFlag + ctypedef struct ZSTD_DCtx: + pass + ctypedef ZSTD_DCtx ZSTD_DStream + + const char* ZSTD_versionString() + unsigned ZSTD_isError(size_t code) + const char* ZSTD_getErrorName(size_t code) + int ZSTD_maxCLevel() + ZSTD_CCtx* ZSTD_createCCtx() + size_t ZSTD_freeCCtx(ZSTD_CCtx* cctx) + ZSTD_CStream* ZSTD_createCStream() + size_t ZSTD_freeCStream(ZSTD_CStream* zcs) + size_t ZSTD_initCStream(ZSTD_CStream* zcs, int compressionLevel) + size_t ZSTD_CCtx_setParameter(ZSTD_CCtx* cctx, ZSTD_cParameter param, int value) + size_t ZSTD_CStreamInSize() + size_t ZSTD_CStreamOutSize() + size_t ZSTD_DStreamInSize() + size_t ZSTD_DStreamOutSize() + size_t ZSTD_compressStream2(ZSTD_CCtx* cctx, ZSTD_outBuffer* output, ZSTD_inBuffer* inp, ZSTD_EndDirective endOp) + ZSTD_DCtx* ZSTD_createDCtx() + size_t ZSTD_freeDCtx(ZSTD_DCtx* dctx) + size_t ZSTD_decompressStream(ZSTD_DStream* zds, ZSTD_outBuffer* output, ZSTD_inBuffer* inp) diff --git a/game_backuper/_zstd.pyx b/game_backuper/_zstd.pyx new file mode 100644 index 0000000..c16e031 --- /dev/null +++ b/game_backuper/_zstd.pyx @@ -0,0 +1,194 @@ +from ._zstd cimport * +from ._Python cimport * +from libc.stdlib cimport malloc, free + + +cdef void CHECK_ZSTD(size_t i) except *: + if ZSTD_isError(i): + raise ValueError(ZSTD_getErrorName(i).decode()) + + +def version(): + cdef const char* v = ZSTD_versionString() + return v.decode() + + +def maxCLevel(): + return ZSTD_maxCLevel() + + +cdef class ZSTDCompressor: + cdef ZSTD_CCtx* cctx + cdef void* buffOut + cdef size_t buffOutSize + cdef int finish + + def __cinit__(self): + self.cctx = NULL + self.buffOut = NULL + self.buffOutSize = 0 + + def __dealloc__(self): + if self.cctx != NULL: + ZSTD_freeCCtx(self.cctx) + if self.buffOut != NULL: + free(self.buffOut) + + def __init__(self, int compresslevel = 3): + if compresslevel < 1 or compresslevel > ZSTD_maxCLevel(): + raise ValueError(u'unsupported compresslevel') + self.finish = 0 + self.cctx = ZSTD_createCCtx() + if self.cctx == NULL: + raise MemoryError() + CHECK_ZSTD(ZSTD_CCtx_setParameter(self.cctx, ZSTD_c_compressionLevel, compresslevel)) + CHECK_ZSTD(ZSTD_CCtx_setParameter(self.cctx, ZSTD_c_checksumFlag, 1)) + + def compress(self, bytes inp): + if self.finish: + raise ValueError('Compressor has been flushed') + if self.buffOut == NULL: + self.buffOutSize = ZSTD_CStreamOutSize() + self.buffOut = malloc(self.buffOutSize) + if self.buffOut == NULL: + raise MemoryError() + cdef int finished = 0 + cdef ZSTD_outBuffer out + cdef ZSTD_inBuffer i + cdef Py_ssize_t si + cdef size_t remaining + cdef char* obuf + if PyBytes_AsStringAndSize(inp, &i.src, &si) == -1: + raise ValueError(u'Can not convert object to void*.') + i.size = si + i.pos = 0 + b = b'' + while not finished: + out.dst = self.buffOut + out.size = self.buffOutSize + out.pos = 0 + remaining = ZSTD_compressStream2(self.cctx, &out, &i, ZSTD_e_continue) + CHECK_ZSTD(remaining) + obuf = out.dst + b += PyBytes_FromStringAndSize(obuf, out.pos) + finished = i.pos == i.size + return b + + def flush(self): + if self.finish: + raise ValueError('Repeated call to flush()') + if self.buffOut == NULL: + self.finish = 1 + return b'' + cdef int finished = 0 + cdef ZSTD_outBuffer out + cdef ZSTD_inBuffer i + cdef char* obuf + i.src = NULL + i.size = 0 + i.pos = 0 + b = b'' + while not finished: + out.dst = self.buffOut + out.size = self.buffOutSize + out.pos = 0 + remaining = ZSTD_compressStream2(self.cctx, &out, &i, ZSTD_e_end) + CHECK_ZSTD(remaining) + obuf = out.dst + b += PyBytes_FromStringAndSize(obuf, out.pos) + finished = remaining == 0 + self.finish = 1 + return b + + +cdef class ZSTDDecompressor: + cdef ZSTD_DCtx* dctx + cdef int finish + cdef object _buff + cdef int need_inp + cdef void* buffOut + cdef size_t buffOutSize + cdef object _unused_data + + def __cinit__(self): + self.dctx = NULL + self._buff = b'' + self._unused_data = b'' + self.buffOut = NULL + self.buffOutSize = 0 + + def __dealloc__(self): + if self.dctx != NULL: + ZSTD_freeDCtx(self.dctx) + if self.buffOut != NULL: + free(self.buffOut) + + def __init__(self): + self.dctx = ZSTD_createDCtx() + if self.dctx == NULL: + raise MemoryError() + self.finish = 0 + self.need_inp = 1 + + def decompress(self, bytes data, Py_ssize_t max_length = -1): + if not self.need_inp: + self.need_inp = 1 + tmp = self._buff + self._buff = b'' + if self.finish: + print(data) + self._unused_data += data + return tmp + if self.finish: + raise EOFError('End of stream already reached') + if self.buffOut == NULL: + self.buffOutSize = ZSTD_DStreamOutSize() + self.buffOut = malloc(self.buffOutSize) + if self.buffOut == NULL: + raise MemoryError() + if self.need_inp: + b = b'' + else: + b = self._buff + self._buff = b'' + cdef int finished = 0 + cdef ZSTD_inBuffer i + cdef ZSTD_outBuffer out + cdef size_t ret + cdef Py_ssize_t si + cdef char* obuf + if PyBytes_AsStringAndSize(data, &i.src, &si) == -1: + raise ValueError(u'Can not convert object to void*.') + i.size = si + i.pos = 0 + while not finished: + out.dst = self.buffOut + out.size = self.buffOutSize + out.pos = 0 + ret = ZSTD_decompressStream(self.dctx, &out, &i) + CHECK_ZSTD(ret) + obuf = out.dst + b += PyBytes_FromStringAndSize(obuf, out.pos) + self.finish = ret == 0 + finished = out.pos < out.size + if self.finish and i.pos < i.size: + obuf = ( i.src) + i.pos + self._unused_data = PyBytes_FromStringAndSize(obuf, i.size - i.pos) + if max_length == 1 or len(b) <= max_length: + return b + else: + self._buff = b[max_length:] + self.need_inp = 0 + return b[:max_length] + + @property + def eof(self): + return True if self.finish else False + + @property + def unused_data(self): + return self._unused_data + + @property + def needs_input(self): + return True if self.need_inp else False diff --git a/game_backuper/compress.py b/game_backuper/compress.py index 8455ffb..5122fd3 100644 --- a/game_backuper/compress.py +++ b/game_backuper/compress.py @@ -102,7 +102,7 @@ class CompressConfig: else: raise ValueError('lzip: compress_level should be 0-9.') self._ext = ".lz" - self._chunk_size = 1048576 + self._chunk_size = 131072 def __repr__(self): t = type(self) diff --git a/setup.py b/setup.py index b364a51..2c1228e 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,15 @@ except ImportError: def cythonize(li): return [] -ext_modules = [Extension("game_backuper._pcre2", ["game_backuper/_pcre2.pyx"], libraries=["pcre2-8"])] +ext_modules = [] +if '--without-pcre2' in sys.argv: + sys.argv.remove('--without-pcre2') +else: + ext_modules.append(Extension("game_backuper._pcre2", ["game_backuper/_pcre2.pyx"], libraries=["pcre2-8"])) +if '--without-zstd' in sys.argv: + sys.argv.remove('--without-zstd') +else: + ext_modules.append(Extension("game_backuper._zstd", ["game_backuper/_zstd.pyx"], libraries=["zstd"])) if "py2exe" in sys.argv: from distutils.core import setup