diff --git a/hypy_utils/__init__.py b/hypy_utils/__init__.py index 8e01d70..c00d5f1 100644 --- a/hypy_utils/__init__.py +++ b/hypy_utils/__init__.py @@ -1,6 +1,6 @@ from __future__ import annotations -__version__ = "1.0.25" +__version__ = "1.0.26" import time import logging diff --git a/hypy_utils/zstd_utils.py b/hypy_utils/zstd_utils.py index 935baa8..391c0c8 100644 --- a/hypy_utils/zstd_utils.py +++ b/hypy_utils/zstd_utils.py @@ -1,3 +1,4 @@ +import pickle from pathlib import Path import zstandard as zstd @@ -5,7 +6,7 @@ import orjson zstd_d = zstd.ZstdDecompressor() -zstd_c = zstd.ZstdCompressor(level=15, write_checksum=True, threads=-1) +zstd_c = zstd.ZstdCompressor(level=5, write_checksum=True, threads=-1) def load_json_zst(file_path: str | Path) -> dict | list: @@ -33,6 +34,17 @@ def write_json_zst(file_path: str | Path, data: dict | list, **kwargs): Path(file_path).write_bytes(zstd_c.compress(orjson.dumps(data, **kwargs))) +def load_pickle_zst(file_path: str | Path) -> object: + with Path(file_path).open('rb') as f: + return pickle.loads(zstd_d.stream_reader(f).read()) + + +def write_pickle_zst(file_path: str | Path, data: object): + Path(file_path).write_bytes(zstd_c.compress(pickle.dumps(data))) + + if __name__ == '__main__': + write_pickle_zst('test.pickle.zst', {'a': 1, 'b': 2}) + assert load_pickle_zst('test.pickle.zst') == {'a': 1, 'b': 2} write_json_zst('test.json.zst', {'a': 1, 'b': 2}) assert load_json_zst('test.json.zst') == {'a': 1, 'b': 2}