Merge pull request #34 from int-brain-lab/jsonable

load jsonable from iblrig available to all installs

CHANGELOG.md

@@ -3,6 +3,13 @@
 * Minor releases (X.1.X) are new features such as added functions or small changes that don't cause major compatibility issues.
 * Major releases (1.X.X) are major new features or changes that break backward compatibility in a big way.
 
+## [Latest](https://github.com/int-brain-lab/iblutil/commits/main) [1.14.0]
+
+### Added
+
+- io.jsonable.load_task_jsonable: read and format iblrig raw data to a trials table Dataframe and a list of raw Bpod trials
+- util.Listable: returns a typing class that is the union of input class and a sequence thereof
+
 ## [Latest](https://github.com/int-brain-lab/iblutil/commits/main) [1.13.0]
 
 ### Added

MANIFEST.ini

@@ -0,0 +1 @@
+recursive-include tests/fixtures *

iblutil/__init__.py

@@ -1 +1 @@
-__version__ = '1.13.0'
+__version__ = '1.14.0'

iblutil/io/jsonable.py

@@ -1,4 +1,7 @@
 import json
+from typing import Any, List, Tuple
+
+import pandas as pd
 
 
 def read(file):
@@ -21,3 +24,37 @@ def write(file, data):
 
 def append(file, data):
     _write(file, data, 'a')
+
+
+def load_task_jsonable(jsonable_file, offset: int = None) -> Tuple[pd.DataFrame, List[Any]]:
+    """
+    Reads in a task data jsonable file and returns a trials dataframe and a bpod data list.
+
+    Parameters
+    ----------
+    jsonable_file : str, pathlib.Path
+        Full path to jsonable file.
+    offset : int
+        The offset to start reading from (default: None).
+
+    Returns
+    -------
+    pandas.DataFrame
+        A DataFrame with the trial info in the same format as the Session trials table.
+    list
+        Timing data for each trial.
+    """
+    trials_table = []
+    with open(jsonable_file) as f:
+        if offset is not None:
+            f.seek(offset, 0)
+        for line in f:
+            trials_table.append(json.loads(line))
+
+    # pop-out the bpod data from the table
+    bpod_data = []
+    for td in trials_table:
+        bpod_data.append(td.pop('behavior_data'))
+
+    trials_table = pd.DataFrame(trials_table)
+    return trials_table, bpod_data

tests/test_io.py → iblutil/tests/test_io.py

@@ -220,5 +220,27 @@ def tearDown(self) -> None:
         os.unlink(self.tfile.name)
 
 
+class TestLoadTaskData(unittest.TestCase):
+    def test_load_task_jsonable(self):
+        jsonable_file = Path(__file__).parent.joinpath('fixtures', 'task_data_short.jsonable')
+        trials_table, bpod_data = jsonable.load_task_jsonable(jsonable_file)
+        assert trials_table.shape[0] == 2
+        assert len(bpod_data) == 2
+
+    def test_load_task_jsonable_partial(self):
+        jsonable_file = Path(__file__).parent.joinpath('fixtures', 'task_data_short.jsonable')
+        with open(jsonable_file) as fp:
+            fp.readline()
+            offset = fp.tell()
+        trials_table, bpod_data = jsonable.load_task_jsonable(jsonable_file, offset=offset)
+
+        trials_table_full, bpod_data_full = jsonable.load_task_jsonable(jsonable_file)
+        for c in trials_table.columns:
+            if not np.isnan(trials_table[c][0]):
+                np.testing.assert_equal(trials_table_full[c].values[-1], trials_table[c][0])
+
+        assert bpod_data_full[-1] == bpod_data[0]
+
+
 if __name__ == '__main__':
     unittest.main(exit=False, verbosity=2)

tests/test_util.py → iblutil/tests/test_util.py

@@ -1,5 +1,6 @@
 import unittest
 import types
+import typing
 from pathlib import Path
 import tempfile
 import logging
@@ -230,5 +231,16 @@ def test_ensure_list(self):
         self.assertEqual([x], util.ensure_list(x, exclude_type=(np.ndarray)))
 
 
+class TestListable(unittest.TestCase):
+    """Test Listable typing class."""
+
+    def test_listable(self):
+        """Test listable type from given class."""
+        listable = util.Listable(str)
+        self.assertIs(listable, typing.Union[str, typing.Sequence[str]])
+        listable = util.Listable(dict)
+        self.assertIs(listable, typing.Union[dict, typing.Sequence[dict]])
+
+
 if __name__ == '__main__':
     unittest.main(exit=False)

iblutil/util.py

@@ -7,7 +7,7 @@
 import copy
 import logging
 import sys
-from typing import Union, Iterable
+from typing import Union, Iterable, Sequence
 
 import numpy as np
 
@@ -23,6 +23,11 @@
     'CRITICAL': 'bold_purple'}
 
 
+def Listable(t):
+    """Return a typing.Union if the input and sequence of input."""
+    return Union[t, Sequence[t]]
+
+
 class Bunch(dict):
     """A subclass of dictionary with an additional dot syntax."""