generalize custom object IO

convnwb/io/io.py

@@ -7,6 +7,7 @@
 
 import yaml
 
+#from convnwb.objects import Task, Bundle, Electrodes
 from convnwb.io.utils import get_files, check_ext, check_folder, make_session_name
 from convnwb.modutils.dependencies import safe_import, check_dependency
 
@@ -142,29 +143,33 @@ def load_configs(files, folder=None):
 
 ### TASK OBJECTS
 
-def save_task_object(task, file_name, folder=None):
-    """Save a task object.
+def save_object(custom_object, file_name, folder=None):
+    """Save a custom object.
 
     Parameters
     ----------
-    task : Task
-        Task object to save out.
+    custom_object : Task or Electrodes
+        Object to save out.
     file_name : str
-        File name to give the saved out task file.
+        File name to give the saved out object.
     folder : str or Path, optional
         Folder to save out to.
 
     Notes
     -----
-    Task objects are saved and loaded as pickle files.
+    Custom objects are saved and loaded as pickle files.
     """
 
-    with open(check_ext(check_folder(file_name, folder), '.task'), 'wb') as fobj:
-        pickle.dump(task, fobj)
+    ext = '.' + str(type(custom_object)).split('.')[-1].strip("'>").lower()
+    if 'task' in ext:
+        ext = '.task'
 
+    with open(check_ext(check_folder(file_name, folder), ext), 'wb') as fobj:
+        pickle.dump(custom_object, fobj)
 
-def load_task_object(file_name, folder=None):
-    """Load a task object.
+
+def load_object(file_name, folder=None):
+    """Load a custom object.
 
     Parameters
     ----------
@@ -175,18 +180,18 @@ def load_task_object(file_name, folder=None):
 
     Returns
     -------
-    task
+    custom_object
         Loaded task object.
 
     Notes
     -----
-    Task objects are saved and loaded as pickle files.
+    Custom objects are saved and loaded as pickle files.
     """
 
-    with open(check_ext(check_folder(file_name, folder), '.task'), 'rb') as load_obj:
-        task = pickle.load(load_obj)
+    with open(check_folder(file_name, folder), 'rb') as load_obj:
+        custom_object = pickle.load(load_obj)
 
-    return task
+    return custom_object
 
 
 ## OTHER FILE I/O

convnwb/tests/io/test_check.py

@@ -5,7 +5,6 @@
 import numpy as np
 import pandas as pd
 
-from convnwb.tests.tobjects import TestTask
 from convnwb.tests.tsettings import TEST_FILE_PATH
 
 from convnwb.io.check import *

convnwb/tests/io/test_io.py

@@ -5,7 +5,6 @@
 import numpy as np
 import pandas as pd
 
-from convnwb.tests.tobjects import TestTask
 from convnwb.tests.tsettings import TEST_FILE_PATH
 
 from convnwb.io.io import *
@@ -51,20 +50,34 @@ def test_load_configs():
     configs = load_configs(f_names, TEST_FILE_PATH)
     assert isinstance(configs, dict)
 
-def test_save_task_object():
+def test_save_object(ttask, tbundle, telectrodes):
 
-    task = TestTask()
     f_name = 'task_obj'
-    save_task_object(task, f_name, TEST_FILE_PATH)
-
+    save_object(ttask, f_name, TEST_FILE_PATH)
     assert os.path.exists(TEST_FILE_PATH / (f_name + '.task'))
 
-def test_load_task_object():
+    f_name = 'bundle_obj'
+    save_object(tbundle, f_name, TEST_FILE_PATH)
+    assert os.path.exists(TEST_FILE_PATH / (f_name + '.bundle'))
 
-    f_name = 'task_obj'
-    task = load_task_object(f_name, TEST_FILE_PATH)
+    f_name = 'electrodes_obj'
+    save_object(telectrodes, f_name, TEST_FILE_PATH)
+    assert os.path.exists(TEST_FILE_PATH / (f_name + '.electrodes'))
+
+def test_load_object():
+
+    f_name = 'task_obj.task'
+    task = load_object(f_name, TEST_FILE_PATH)
     assert task
 
+    f_name = 'bundle_obj.bundle'
+    bundle = load_object(f_name, TEST_FILE_PATH)
+    assert bundle
+
+    f_name = 'electrodes_obj.electrodes'
+    electrodes = load_object(f_name, TEST_FILE_PATH)
+    assert electrodes
+
 def test_save_txt():
 
     text = "Words, words, words."

doc/api.rst

@@ -30,8 +30,8 @@ NWB & Custom file I/O
    save_config
    load_config
    load_configs
-   save_task_object
-   load_task_object
+   save_object
+   load_object
 
 General file I/O
 ~~~~~~~~~~~~~~~~