diff --git a/hypy_utils/__init__.py b/hypy_utils/__init__.py index c00d5f1..fb05f1e 100644 --- a/hypy_utils/__init__.py +++ b/hypy_utils/__init__.py @@ -1,6 +1,6 @@ from __future__ import annotations -__version__ = "1.0.26" +__version__ = "1.0.27" import time import logging diff --git a/hypy_utils/zstd_utils.py b/hypy_utils/zstd_utils.py index 391c0c8..dd99609 100644 --- a/hypy_utils/zstd_utils.py +++ b/hypy_utils/zstd_utils.py @@ -4,6 +4,7 @@ from pathlib import Path import zstandard as zstd import orjson +from . import write zstd_d = zstd.ZstdDecompressor() zstd_c = zstd.ZstdCompressor(level=5, write_checksum=True, threads=-1) @@ -31,16 +32,16 @@ def write_json_zst(file_path: str | Path, data: dict | list, **kwargs): file_path (str): The path to the .json.zst file. data (dict or list): The data to dump. """ - Path(file_path).write_bytes(zstd_c.compress(orjson.dumps(data, **kwargs))) + write(file_path, zstd_c.compress(orjson.dumps(data, **kwargs))) -def load_pickle_zst(file_path: str | Path) -> object: +def load_pickle_zst(file_path: str | Path): 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))) +def write_pickle_zst(file_path: str | Path, data): + write(file_path, zstd_c.compress(pickle.dumps(data))) if __name__ == '__main__':