NeoMatlabIO: round-trip quantity, array and nested annotations

doc/source/authors.rst

@@ -103,6 +103,7 @@ and may not be the current affiliation of a contributor.
 * Reema Gupta [50]
 * Sai Asish Yamani [51]
 * Kevin Doran [52]
+* Aditya Singh (github)
 
 1. Centre de Recherche en Neuroscience de Lyon, CNRS UMR5292 - INSERM U1028 - Université Claude Bernard Lyon 1
 2. Unité de Neuroscience, Information et Complexité, CNRS UPR 3293, Gif-sur-Yvette, France

neo/io/neomatlabio.py

@@ -334,11 +334,13 @@ def _get_matlab_value(self, ob, attrname):
                         new_value[key] = subvalue
                         if subunits:
                             new_value[f"{key}_units"] = subunits
-                    elif attrname == "annotations":
+                    else:
                         # In general we don't send None to MATLAB
                         # but we make an exception for annotations.
                         # However, we have to save then retrieve some
                         # special value as actual `None` is ignored by default.
+                        # Annotations are the only mapping-valued attribute Neo has,
+                        # so this holds at every level of a nested annotation too.
                         new_value[key] = PY_NONE
                 value = new_value
         return value, units
@@ -377,6 +379,37 @@ def create_struct_from_view(self, ob):
         struct["viewed_classname"] = viewed_obj.__class__.__name__
         return struct
 
+    def create_dict_from_struct(self, struct):
+        """
+        Rebuild a Python dict, typically the annotations, from the MATLAB struct that
+        :meth:`_get_matlab_value` wrote it to.
+
+        That method flattens a quantity into a plain magnitude plus a companion
+        ``<key>_units`` field, mirrors nested mappings as nested structs, and stores
+        `None` as a sentinel string because MATLAB has no equivalent. This undoes all
+        three so that a value survives a write/read round trip unchanged.
+        """
+        new_dict = {}
+        for field_name in struct._fieldnames:
+            if field_name.endswith("_units") and field_name[: -len("_units")] in struct._fieldnames:
+                # this field carries the units of another one and is consumed along with it
+                continue
+
+            value = getattr(struct, field_name)
+            if hasattr(value, "_fieldnames"):
+                # a nested mapping, which scipy returns as a struct of its own
+                value = self.create_dict_from_struct(value)
+            elif isinstance(value, str) and value == PY_NONE:
+                # `isinstance` first: an array compared to the sentinel gives an array,
+                # which is not usable as a condition
+                value = None
+            else:
+                units = getattr(struct, f"{field_name}_units", None)
+                if units is not None:
+                    value = pq.Quantity(value, str(units))
+            new_dict[field_name] = value
+        return new_dict
+
     def create_ob_from_struct(self, struct, classname):
         cl = class_by_name[classname]
 
@@ -508,13 +541,7 @@ def create_ob_from_struct(self, struct, classname):
                     else:
                         item = pq.Quantity(item, units)
                 elif attrtype == dict:
-                    new_item = {}
-                    for fn in item._fieldnames:
-                        value = getattr(item, fn)
-                        if value == PY_NONE:
-                            value = None
-                        new_item[fn] = value
-                    item = new_item
+                    item = self.create_dict_from_struct(item)
                 else:
                     item = attrtype(item)
 

neo/test/iotest/test_neomatlabio.py

@@ -4,6 +4,7 @@
 
 import os
 import unittest
+import pytest
 from numpy.testing import assert_array_equal
 import quantities as pq
 
@@ -33,7 +34,7 @@ def test_write_read_single_spike(self):
         block1 = Block(name="test_neomatlabio")
         seg = Segment("segment1")
         spiketrain1 = SpikeTrain([1] * pq.s, t_stop=10 * pq.s, sampling_rate=1 * pq.Hz)
-        spiketrain1.annotate(yep="yop", yip=None)
+        spiketrain1.annotate(yep="yop", yip=None, yop=[3, 4, 5] * pq.ms)
         sig1 = AnalogSignal([4, 5, 6] * pq.A, sampling_period=1 * pq.ms)
         irrsig1 = IrregularlySampledSignal([0, 1, 2] * pq.ms, [4, 5, 6] * pq.A)
         img_sequence_array = [[[column for column in range(2)] for _ in range(2)] for _ in range(2)]
@@ -76,6 +77,11 @@ def test_write_read_single_spike(self):
         spiketrain2 = block2.segments[0].spiketrains[0]
         assert spiketrain2.annotations["yep"] == "yop"
         assert spiketrain2.annotations["yip"] is None
+        # a quantity annotation keeps its units and does not leak its companion field,
+        # see https://github.com/NeuralEnsemble/python-neo/issues/852
+        assert "yop_units" not in spiketrain2.annotations
+        assert spiketrain2.annotations["yop"].dimensionality == pq.ms.dimensionality
+        assert_array_equal(spiketrain2.annotations["yop"].magnitude, [3, 4, 5])
 
         # test group retrieval
         group2 = block2.groups[0]
@@ -104,5 +110,65 @@ def test_write_read_random_blocks(self):
             assert os.stat(filename_orig).st_size == os.stat(filename_roundtripped).st_size
 
 
+@unittest.skipUnless(HAVE_SCIPY, "requires scipy")
+@pytest.mark.parametrize(
+    "annotation",
+    [
+        "a string",
+        42,
+        None,
+        [3, 4, 5],
+        [3, 4, 5] * pq.ms,
+        3 * pq.ms,
+        {"a": 1, "b": None},
+        {"a": [1, 2] * pq.mV, "b": {"deep": 5 * pq.Hz}},
+    ],
+    ids=["str", "int", "none", "array", "quantity_array", "quantity_scalar", "dict", "nested_dict"],
+)
+def test_write_read_annotation_roundtrip(tmp_path, annotation):
+    """An annotation must survive a write then read unchanged.
+
+    A test case from Issue #852 (https://github.com/NeuralEnsemble/python-neo/issues/852).
+    Quantity annotations used to come back as bare magnitudes plus a stray ``<key>_units``
+    key, nested mappings came back as ``scipy.io`` structs, and any array-valued
+    annotation raised ``ValueError: The truth value of an array ... is ambiguous`` while
+    being compared against the sentinel that stands in for `None`.
+
+    This runs off a temporary file rather than the downloaded test data, so it also
+    covers installations without datalad.
+    """
+    block = Block(name="test_annotations")
+    segment = Segment(name="segment1")
+    block.segments.append(segment)
+    spiketrain = SpikeTrain([1, 2, 3] * pq.s, t_stop=10 * pq.s, yop=annotation)
+    segment.spiketrains.append(spiketrain)
+    segment.check_relationships()
+
+    filename = tmp_path / "annotations.mat"
+    NeoMatlabIO(filename=filename).write_block(block)
+    read_back = NeoMatlabIO(filename=filename).read_block().segments[0].spiketrains[0].annotations
+
+    assert list(read_back) == ["yop"], "no companion field should leak into the annotations"
+    _assert_annotation_equal(read_back["yop"], annotation)
+
+
+def _assert_annotation_equal(actual, expected):
+    if expected is None:
+        assert actual is None
+    elif isinstance(expected, dict):
+        assert isinstance(actual, dict)
+        assert sorted(actual) == sorted(expected)
+        for key in expected:
+            _assert_annotation_equal(actual[key], expected[key])
+    elif isinstance(expected, pq.Quantity):
+        assert isinstance(actual, pq.Quantity)
+        assert actual.dimensionality == expected.dimensionality
+        assert_array_equal(actual.magnitude, expected.magnitude)
+    elif isinstance(expected, list):
+        assert_array_equal(actual, expected)
+    else:
+        assert actual == expected
+
+
 if __name__ == "__main__":
     unittest.main()