Modify spec, add tests

spec/ndx-hed.extensions.yaml

@@ -1,8 +1,15 @@
+datasets:
+- neurodata_type_def: HedTags
+  neurodata_type_inc: VectorData
+  doc: An extension of VectorData for Hierarchical Event Descriptor (HED) tags.
 groups:
-- neurodata_type_def: TetrodeSeries
-  neurodata_type_inc: ElectricalSeries
-  doc: An extension of ElectricalSeries to include the tetrode ID for each time series.
+- neurodata_type_def: HedNWBFile
+  neurodata_type_inc: NWBFile
+  doc: An extension of NWBFile to store the Hierarchical Event Descriptor (HED) schema
+    version.
   attributes:
-  - name: trode_id
-    dtype: int32
-    doc: The tetrode ID.
+  - name: hed_schema_version
+    dtype: text
+    doc: The version of the HED schema used to validate the HED tags, e.g., '8.2.0'.
+      Required if HED tags are used in the NWB file.
+    required: false

src/pynwb/ndx_hed/__init__.py

@@ -22,7 +22,8 @@
 # Either have PyNWB generate a class from the spec using `get_class` as shown
 # below or write a custom class and register it using the class decorator
 # `@register_class("TetrodeSeries", "ndx-hed")`
-TetrodeSeries = get_class("TetrodeSeries", "ndx-hed")
+HedTags = get_class("HedTags", "ndx-hed")
+HedNWBFile = get_class("HedNWBFile", "ndx-hed")
 
 # Remove these functions from the package
 del load_namespaces, get_class

src/pynwb/tests/test_hed_tags.py

@@ -0,0 +1,164 @@
+"""Unit and integration tests for ndx-hed."""
+
+import numpy as np
+
+from pynwb import NWBHDF5IO, NWBFile
+from pynwb.testing.mock.device import mock_Device
+from pynwb.testing.mock.ecephys import mock_ElectrodeGroup, mock_ElectrodeTable
+from pynwb.testing.mock.file import mock_NWBFile
+from pynwb.testing import TestCase, remove_test_file, NWBH5IOFlexMixin
+
+from hdmf.common import VectorIndex
+from uuid import uuid4
+from datetime import datetime
+from dateutil.tz import tzlocal
+
+from ndx_hed import HedTags, HedNWBFile
+
+
+class TestHedNWBFileConstructor(TestCase):
+    """Simple unit test for creating a HedNWBFile."""
+
+    def test_constructor(self):
+        """Test setting HedNWBFile values using the constructor."""
+        hed_nwbfile = HedNWBFile(
+            session_description="session_description",
+            identifier=str(uuid4()),
+            session_start_time=datetime(1970, 1, 1, tzinfo=tzlocal()),
+            hed_schema_version="8.2.0",
+        )
+        assert hed_nwbfile.hed_schema_version == "8.2.0"
+
+
+class TestHedNWBFileSimpleRoundtrip(TestCase):
+    """Simple roundtrip test for HedNWBFile."""
+
+    def setUp(self):
+        self.path = "test.nwb"
+
+    def tearDown(self):
+        remove_test_file(self.path)
+
+    def test_roundtrip(self):
+        """
+        Create a HedNWBFile, write it to file, read the file, and test that it matches the original HedNWBFile.
+        """
+        hed_nwbfile = HedNWBFile(
+            session_description="session_description",
+            identifier=str(uuid4()),
+            session_start_time=datetime(1970, 1, 1, tzinfo=tzlocal()),
+            hed_schema_version="8.2.0",
+        )
+
+        with NWBHDF5IO(self.path, mode="w") as io:
+            io.write(hed_nwbfile)
+
+        with NWBHDF5IO(self.path, mode="r", load_namespaces=True) as io:
+            read_nwbfile = io.read()
+            assert isinstance(read_nwbfile, HedNWBFile)
+            assert read_nwbfile.hed_schema_version == "8.2.0"
+
+
+class TestHedNWBFileRoundtripPyNWB(NWBH5IOFlexMixin, TestCase):
+    """Complex, more complete roundtrip test for HedNWBFile using pynwb.testing infrastructure."""
+
+    def getContainerType(self):
+        return "HedNWBFile"
+
+    def addContainer(self):
+        self.nwbfile = HedNWBFile(
+            session_description="session_description",
+            identifier=str(uuid4()),
+            session_start_time=datetime(1970, 1, 1, tzinfo=tzlocal()),
+            hed_schema_version="8.2.0",
+        )
+
+    def getContainer(self, nwbfile: NWBFile):
+        return nwbfile
+
+
+class TestHedTagsConstructor(TestCase):
+    """Simple unit test for creating a HedTags."""
+
+    def test_constructor(self):
+        """Test setting HedTags values using the constructor."""
+        hed_tags = HedTags(
+            name="name",
+            description="description",
+            data=["tag1", "tag2"],
+        )
+        assert hed_tags.name == "name"
+        assert hed_tags.description == "description"
+        assert hed_tags.data == ["tag1", "tag2"]
+
+    def test_add_to_trials_table(self):
+        """Test adding HedTags column and data to a trials table."""
+        nwbfile = mock_NWBFile()
+        nwbfile.add_trial_column("hed_tags", "HED tags for each trial", col_cls=HedTags, index=True)
+        nwbfile.add_trial(start_time=0.0, stop_time=1.0, hed_tags=["tag1", "tag2"])
+        nwbfile.add_trial(start_time=2.0, stop_time=3.0, hed_tags=["tag1", "tag3"])
+
+        assert isinstance(nwbfile.trials["hed_tags"], VectorIndex)
+        assert isinstance(nwbfile.trials["hed_tags"].target, HedTags)
+        assert nwbfile.trials["hed_tags"][0] == ["tag1", "tag2"]
+        assert nwbfile.trials["hed_tags"][0] == ["tag1", "tag2"]
+
+
+class TestHedTagsSimpleRoundtrip(TestCase):
+    """Simple roundtrip test for HedTags."""
+
+    def setUp(self):
+        self.path = "test.nwb"
+
+    def tearDown(self):
+        remove_test_file(self.path)
+
+    def test_roundtrip(self):
+        """
+        Add a HedTags to an NWBFile, write it to file, read the file, and test that the HedTags from the
+        file matches the original HedTags.
+        """
+        hed_nwbfile = HedNWBFile(
+            session_description="session_description",
+            identifier=str(uuid4()),
+            session_start_time=datetime(1970, 1, 1, tzinfo=tzlocal()),
+            hed_schema_version="8.2.0",
+        )
+
+        hed_nwbfile.add_trial_column("hed_tags", "HED tags for each trial", col_cls=HedTags, index=True)
+        hed_nwbfile.add_trial(start_time=0.0, stop_time=1.0, hed_tags=["tag1", "tag2"])
+        hed_nwbfile.add_trial(start_time=2.0, stop_time=3.0, hed_tags=["tag1", "tag3"])
+
+        with NWBHDF5IO(self.path, mode="w") as io:
+            io.write(hed_nwbfile)
+
+        with NWBHDF5IO(self.path, mode="r", load_namespaces=True) as io:
+            read_nwbfile = io.read()
+            assert isinstance(read_nwbfile, HedNWBFile)
+            assert isinstance(read_nwbfile.trials["hed_tags"], VectorIndex)
+            assert isinstance(read_nwbfile.trials["hed_tags"].target, HedTags)
+            # read_nwbfile.trials["hed_tags"][0] is read as a numpy array
+            assert all(read_nwbfile.trials["hed_tags"][0] == ["tag1", "tag2"])
+            assert all(read_nwbfile.trials["hed_tags"][1] == ["tag1", "tag3"])
+
+
+class TestHedTagsRoundtripPyNWB(NWBH5IOFlexMixin, TestCase):
+    """Complex, more complete roundtrip test for HedTags using pynwb.testing infrastructure."""
+
+    def getContainerType(self):
+        return "HedTags"
+
+    def addContainer(self):
+        self.nwbfile = HedNWBFile(
+            session_description="session_description",
+            identifier=str(uuid4()),
+            session_start_time=datetime(1970, 1, 1, tzinfo=tzlocal()),
+            hed_schema_version="8.2.0",
+        )
+
+        self.nwbfile.add_trial_column("hed_tags", "HED tags for each trial", col_cls=HedTags, index=True)
+        self.nwbfile.add_trial(start_time=0.0, stop_time=1.0, hed_tags=["tag1", "tag2"])
+        self.nwbfile.add_trial(start_time=2.0, stop_time=3.0, hed_tags=["tag1", "tag3"])
+
+    def getContainer(self, nwbfile: NWBFile):
+        return nwbfile.trials["hed_tags"].target

src/spec/create_extension_spec.py

@@ -1,10 +1,10 @@
 # -*- coding: utf-8 -*-
 import os.path
 
-from pynwb.spec import NWBNamespaceBuilder, export_spec, NWBGroupSpec, NWBAttributeSpec
+from pynwb.spec import NWBNamespaceBuilder, export_spec, NWBDatasetSpec, NWBGroupSpec, NWBAttributeSpec
 
 # TODO: import other spec classes as needed
-# from pynwb.spec import NWBDatasetSpec, NWBLinkSpec, NWBDtypeSpec, NWBRefSpec
+# from pynwb.spec import , NWBLinkSpec, NWBDtypeSpec, NWBRefSpec
 
 
 def main():
@@ -21,7 +21,7 @@ def main():
         contact=[
             "[email protected]",
             "[email protected]",
-            "[email protected]"
+            "[email protected]",
         ],
     )
 
@@ -36,15 +36,31 @@ def main():
     # TODO: define your new data types
     # see https://pynwb.readthedocs.io/en/latest/extensions.html#extending-nwb
     # for more information
-    tetrode_series = NWBGroupSpec(
-        neurodata_type_def="TetrodeSeries",
-        neurodata_type_inc="ElectricalSeries",
-        doc="An extension of ElectricalSeries to include the tetrode ID for each time series.",
-        attributes=[NWBAttributeSpec(name="trode_id", doc="The tetrode ID.", dtype="int32")],
+    hed_tags = NWBDatasetSpec(
+        neurodata_type_def="HedTags",
+        neurodata_type_inc="VectorData",
+        doc="An extension of VectorData for Hierarchical Event Descriptor (HED) tags.",
+    )
+
+    hed_nwbfile = NWBGroupSpec(
+        neurodata_type_def="HedNWBFile",
+        neurodata_type_inc="NWBFile",
+        doc="An extension of NWBFile to store the Hierarchical Event Descriptor (HED) schema version.",
+        attributes=[
+            NWBAttributeSpec(
+                name="hed_schema_version",
+                doc=(
+                    "The version of the HED schema used to validate the HED tags, e.g., '8.2.0'. "
+                    "Required if HED tags are used in the NWB file."
+                ),
+                dtype="text",
+                required=False,
+            ),
+        ],
     )
 
     # TODO: add all of your new data types to this list
-    new_data_types = [tetrode_series]
+    new_data_types = [hed_tags, hed_nwbfile]
 
     # export the spec to yaml files in the spec folder
     output_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "spec"))