feat: joblib driver support

Author: machow

pins/drivers.py

@@ -4,19 +4,37 @@
 # from .boards import IFileSystem
 
 
-def load_data(meta: Meta, fs, path_to_version):
+def load_data(meta: Meta, fs, path_to_version: str):
+    """Return loaded data, based on meta type.
+    Parameters
+    ----------
+    meta: Meta
+        Information about the stored data (e.g. its type).
+    fs: IFileSystem
+        An abstract filesystem with a method to .open() files.
+    path_to_version:
+        A filepath used as the parent directory the data to-be-loaded lives in.
+    """
     # TODO: extandable loading with deferred importing
+
+    # Check that only a single file name was given
+    fnames = [meta.file] if isinstance(meta.file, str) else meta.file
+    if len(fnames) > 1:
+        raise ValueError("Cannot load data when more than 1 file")
+
+    target_fname = fnames[0]
+    path_to_file = f"{path_to_version}/{target_fname}"
+
     if meta.type == "csv":
         import pandas as pd
 
-        fnames = [meta.file] if isinstance(meta.file, str) else meta.file
-        if len(fnames) > 1:
-            raise ValueError("Cannot load CSV when more than 1 file")
-
-        target_fname = fnames[0]
-        path_to_file = f"{path_to_version}/{target_fname}"
         return pd.read_csv(fs.open(path_to_file), index_col=0)
 
+    elif meta.type == "joblib":
+        import joblib
+
+        return joblib.load(fs.open(path_to_file))
+
     raise NotImplementedError(f"No driver for type {meta.type}")
 
 
@@ -35,6 +53,10 @@ def save_data(obj, fname, type=None):
                 "Currently only pandas.DataFrame can be saved to a CSV."
             )
         obj.to_csv(fname)
+    elif type == "joblib":
+        import joblib
+
+        joblib.dump(obj, fname)
     else:
         raise NotImplementedError(f"Cannot save type: {type}")
 

pins/tests/test_boards.py

@@ -1,24 +1,34 @@
 import pytest
+import pandas as pd
 
 from pins.tests.helpers import DEFAULT_CREATION_DATE
 
+# using pytest cases, so that we can pass in fixtures as parameters
+from pytest_cases import fixture, parametrize
 
-@pytest.fixture
+
+@fixture
 def board(backend):
     yield backend.create_tmp_board()
     backend.teardown()
 
 
+@fixture
+def df():
+    return pd.DataFrame({"x": [1, 2, 3], "y": ["a", "b", "c"]})
+
+
+# High level pins functionality -----------------------------------------------
+
+
 def test_board_pin_write_default_title(board):
-    import pandas as pd
 
     df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
     meta = board.pin_write(df, "df_csv", title=None, type="csv")
     assert meta.title == "A pinned 3 x 2 CSV"
 
 
 def test_board_pin_write_prepare_pin(board, tmp_dir2):
-    import pandas as pd
 
     df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
 
@@ -32,7 +42,6 @@ def test_board_pin_write_prepare_pin(board, tmp_dir2):
 
 
 def test_board_pin_write_roundtrip(board):
-    import pandas as pd
 
     df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
 
@@ -68,8 +77,6 @@ def test_board_pin_write_rsc_index_html(board, tmp_dir2, snapshot):
     if board.fs.protocol != "rsc":
         pytest.skip()
 
-    import pandas as pd
-
     df = pd.DataFrame({"x": [1, 2, 3], "y": ["a", "b", "c"]})
 
     pin_name = "test_rsc_pin"
@@ -85,3 +92,21 @@ def test_board_pin_write_rsc_index_html(board, tmp_dir2, snapshot):
     )
 
     snapshot.assert_equal_dir(tmp_dir2)
+
+
+# pin_write against different types -------------------------------------------
+
+
+@parametrize(
+    "obj, type_", [(df, "csv"), (df, "joblib"), ({"a": 1, "b": [2, 3]}, "joblib")]
+)
+def test_board_pin_write_type(board, obj, type_, request):
+    meta = board.pin_write(obj, "test_pin", type=type_, title="some title")
+    dst_obj = board.pin_read("test_pin")
+
+    assert meta.type == type_
+
+    if isinstance(obj, pd.DataFrame):
+        assert obj.equals(dst_obj)
+
+    obj == dst_obj