feat: refactor pypots.data.saving package and add testing cases;

pypots/data/saving/__init__.py

@@ -0,0 +1,15 @@
+"""
+Data saving utilities.
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+from .h5 import save_dict_into_h5
+from .pickle import pickle_dump, pickle_load
+
+__all__ = [
+    "save_dict_into_h5",
+    "pickle_dump",
+    "pickle_load",
+]

pypots/data/saving.py → pypots/data/saving/h5.py

@@ -1,5 +1,5 @@
 """
-Data saving utilities.
+Data saving utilities with HDF5.
 """
 
 # Created by Wenjie Du <wenjay.du@gmail.com>
@@ -9,11 +9,10 @@
 import os
 from typing import Optional
 
-import dill as pickle
 import h5py
 
-from ..utils.file import extract_parent_dir, create_dir_if_not_exist
-from ..utils.logging import logger
+from pypots.utils.file import extract_parent_dir, create_dir_if_not_exist
+from pypots.utils.logging import logger
 
 
 def save_dict_into_h5(
@@ -85,53 +84,3 @@ def save_set(handle, name, data):
             save_set(hf, k, v)
 
     logger.info(f"Successfully saved the given data into {saving_path}.")
-
-
-def pickle_dump(data: object, path: str) -> Optional[str]:
-    """Pickle the given object.
-
-    Parameters
-    ----------
-    data:
-        The object to be pickled.
-
-    path:
-        Saving path.
-
-    Returns
-    -------
-    `path` if succeed else None
-
-    """
-    try:
-        with open(path, "wb") as f:
-            pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
-    except pickle.PicklingError:
-        logger.error("❌ Pickling failed. No cache data saved.")
-        return None
-    logger.info(f"Successfully saved to {path}")
-    return path
-
-
-def pickle_load(path: str) -> object:
-    """Load pickled object from file.
-
-    Parameters
-    ----------
-    path :
-        Local path of the pickled object.
-
-    Returns
-    -------
-    Object
-        Pickled object.
-
-    """
-    try:
-        with open(path, "rb") as f:
-            data = pickle.load(f)
-    except pickle.UnpicklingError as e:
-        logger.error(
-            "❌ Data file corrupted. Operation aborted. See info below:\n" f"{e}"
-        )
-    return data

pypots/data/saving/pickle.py

@@ -0,0 +1,61 @@
+"""
+Data saving utilities with pickle.
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+import pickle
+from typing import Optional
+
+from pypots.utils.logging import logger
+
+
+def pickle_dump(data: object, path: str) -> Optional[str]:
+    """Pickle the given object.
+
+    Parameters
+    ----------
+    data:
+        The object to be pickled.
+
+    path:
+        Saving path.
+
+    Returns
+    -------
+    `path` if succeed else None
+
+    """
+    try:
+        with open(path, "wb") as f:
+            pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
+    except pickle.PicklingError:
+        logger.error("❌ Pickling failed. No cache data saved.")
+        return None
+    logger.info(f"Successfully saved to {path}")
+    return path
+
+
+def pickle_load(path: str) -> object:
+    """Load pickled object from file.
+
+    Parameters
+    ----------
+    path :
+        Local path of the pickled object.
+
+    Returns
+    -------
+    Object
+        Pickled object.
+
+    """
+    try:
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+    except pickle.UnpicklingError as e:
+        logger.error(
+            "❌ Data file corrupted. Operation aborted. See info below:\n" f"{e}"
+        )
+    return data

tests/data/saving.py

@@ -0,0 +1,36 @@
+"""
+Test cases for data saving utils.
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+
+import unittest
+
+import pytest
+
+from pypots.data.saving import save_dict_into_h5, pickle_dump, pickle_load
+from pypots.utils.logging import logger
+
+
+class TestLazyLoadingClasses(unittest.TestCase):
+    logger.info("Running tests for data saving utils...")
+
+    data_to_save = {
+        "a": 1,
+        "b": 2,
+        "c": {
+            "d": 0,
+        },
+    }
+
+    @pytest.mark.xdist_group(name="data-saving-h5")
+    def test_0_save_dict_into_h5(self):
+        save_dict_into_h5(self.data_to_save, "tests/data/saving_with_h5.h5")
+
+    @pytest.mark.xdist_group(name="data-saving-pickle")
+    def test_0_pickle_dump_load(self):
+        pickle_dump(self.data_to_save, "tests/data/saving_with_pickle.pkl")
+        loaded_data = pickle_load("tests/data/saving_with_pickle.pkl")
+        assert loaded_data == self.data_to_save