From 27e3f92186f2324e6f090dd43333d2afa14a56ae Mon Sep 17 00:00:00 2001
From: Azalea Gui <22280294+hykilpikonna@users.noreply.github.com>
Date: Thu, 28 Nov 2024 01:03:46 -0500
Subject: [PATCH] [+] Pickle zst

---
 hypy_utils/__init__.py   |  2 +-
 hypy_utils/zstd_utils.py | 14 +++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

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}