TermSetWrapper and write support

docs/write_foo.py

@@ -0,0 +1,52 @@
+from datetime import datetime
+from uuid import uuid4
+
+import numpy as np
+from dateutil.tz import tzlocal
+
+from pynwb import NWBHDF5IO, NWBFile
+from pynwb.ecephys import LFP, ElectricalSeries
+
+from hdmf import TermSetWrapper as tw
+from hdmf import Data
+from hdmf import TermSet
+terms = TermSet(term_schema_path='/Users/mavaylon/Research/NWB/hdmf2/hdmf/docs/gallery/example_term_set.yaml')
+
+import numpy as np
+
+from pynwb import TimeSeries
+
+data = np.arange(100, 200, 10)
+timestamps = np.arange(10)
+
+from hdmf.backends.hdf5.h5_utils import H5DataIO
+
+test_ts = TimeSeries(
+    name="test_compressed_timeseries",
+    data=H5DataIO(data=data, compression=True),
+    unit=tw(item="SIunit", termset=terms),
+    timestamps=timestamps,
+)
+
+nwbfile = NWBFile(
+    session_description="my first synthetic recording",
+    identifier=str(uuid4()),
+    session_start_time=datetime.now(tzlocal()),
+    experimenter=[
+        "Baggins, Bilbo",
+    ],
+    lab="Bag End Laboratory",
+    institution="University of Middle Earth at the Shire",
+    experiment_description="I went on an adventure to reclaim vast treasures.",
+    session_id="LONELYMTN001",
+)
+nwbfile.add_acquisition(test_ts)
+
+filename = "nwbfile_test.nwb"
+with NWBHDF5IO(filename, "w") as io:
+    io.write(nwbfile)
+
+# open the NWB file in r+ mode
+with NWBHDF5IO(filename, "r+") as io:
+    read_nwbfile = io.read()
+breakpoint()

src/hdmf/__init__.py

@@ -1,6 +1,6 @@
 from . import query
 from .backends.hdf5.h5_utils import H5Dataset, H5RegionSlicer
-from .container import Container, Data, DataRegion, HERDManager
+from .container import Container, Data, DataRegion, HERDManager, TermSetWrapper
 from .region import ListSlicer
 from .utils import docval, getargs
 from .term_set import TermSet

src/hdmf/backends/hdf5/h5tools.py

@@ -16,7 +16,7 @@
 from ..warnings import BrokenLinkWarning
 from ...build import (Builder, GroupBuilder, DatasetBuilder, LinkBuilder, BuildManager, RegionBuilder,
                       ReferenceBuilder, TypeMap, ObjectMapper)
-from ...container import Container
+from ...container import Container, TermSetWrapper
 from ...data_utils import AbstractDataChunkIterator
 from ...spec import RefSpec, DtypeSpec, NamespaceCatalog
 from ...utils import docval, getargs, popargs, get_data_shape, get_docval, StrDataset
@@ -1099,6 +1099,8 @@ def write_dataset(self, **kwargs):  # noqa: C901
             dataio = data
             link_data = data.link_data
             data = data.data
+        if isinstance(data, TermSetWrapper):
+            data = data.item
         else:
             options['io_settings'] = {}
         attributes = builder.attributes

src/hdmf/build/objectmapper.py

@@ -11,7 +11,7 @@
                      ConstructError)
 from .manager import Proxy, BuildManager
 from .warnings import MissingRequiredBuildWarning, DtypeConversionWarning, IncorrectQuantityBuildWarning
-from ..container import AbstractContainer, Data, DataRegion
+from ..container import AbstractContainer, Data, DataRegion, TermSetWrapper
 from ..data_utils import DataIO, AbstractDataChunkIterator
 from ..query import ReferenceResolver
 from ..spec import Spec, AttributeSpec, DatasetSpec, GroupSpec, LinkSpec, RefSpec
@@ -564,6 +564,8 @@ def get_attr_value(self, **kwargs):
                 msg = ("%s '%s' does not have attribute '%s' for mapping to spec: %s"
                        % (container.__class__.__name__, container.name, attr_name, spec))
                 raise ContainerConfigurationError(msg)
+            if isinstance(attr_val, TermSetWrapper):
+                    attr_val = attr_val.item
             if attr_val is not None:
                 attr_val = self.__convert_string(attr_val, spec)
                 spec_dt = self.__get_data_type(spec)
@@ -906,7 +908,7 @@ def __add_attributes(self, builder, attributes, container, build_manager, source
             if spec.value is not None:
                 attr_value = spec.value
             else:
-                attr_value = self.get_attr_value(spec, container, build_manager)
+                attr_value = self.get_attr_value(spec, container, build_manager) # here
                 if attr_value is None:
                     attr_value = spec.default_value
 
@@ -937,7 +939,6 @@ def __add_attributes(self, builder, attributes, container, build_manager, source
                 if attr_value is None:
                     self.logger.debug("        Skipping empty attribute")
                     continue
-
             builder.set_attribute(spec.name, attr_value)
 
     def __set_attr_to_ref(self, builder, attr_value, build_manager, spec):

src/hdmf/container.py

@@ -10,7 +10,7 @@
 import pandas as pd
 
 from .data_utils import DataIO, append_data, extend_data
-from .utils import docval, get_docval, getargs, ExtenderMeta, get_data_shape, popargs, LabelledDict
+from .utils import docval, docval_macro, get_docval, getargs, ExtenderMeta, get_data_shape, popargs, LabelledDict
 from hdmf.term_set import TermSet
 
 
@@ -1485,3 +1485,78 @@
         if name is None:
             return cls(data=data)
         return cls(name=name, data=data)
+
+
+class TermSetWrapper:
+    """
+    This class allows any HDF5 group, dataset, or attribute to have a TermSet.
+
+    In HDMF, a group is a Container object, a dataset is a Data object,
+    an attribute can be a reference type to an HDMF object or a base type, e.g., text.
+    """
+    # @docval({'name': 'termset',
+    #          'type': TermSet,
+    #          'doc': 'The TermSet to be used.'},
+    #         {'name': primitive})
+    def __init__(self, **kwargs):
+        item, termset = popargs('item', 'termset', kwargs)
+
+        self.__item = item
+        self.__termset = termset
+        self.__validate()
+
+    def __validate(self):
+        # check if list, tuple, array, or DataIO
+        if isinstance(self.__item, (list, np.ndarray, tuple, Data, DataIO, dict)):
+            values = self.__item
+        # create list if none of those
+        else:
+            values = [self.__item]
+        # iteratively validate
+        bad_values = []
+        for term in values:
+            validation = self.__termset.validate(term=term)
+            if not validation:
+                bad_values.append(term)
+        if len(bad_values)!=0:
+            msg = ('"%s" is not in the term set.' % ', '.join([str(item) for item in bad_values]))
+            raise ValueError(msg)
+
+    @property
+    def item(self):
+        return self.__item
+
+    @property
+    def termset(self):
+        return self.__termset
+
+    @property
+    def dtype(self):
+        return self.__getattr__('dtype')
+
+    def __getattr__(self, val):
+        """
+        This method is to get attributes that are not defined in init.
+        This is when dealing with data and numpy arrays.
+        """
+        if val in ('data', 'shape', 'dtype'):
+            return getattr(self.__item, val)
+
+    def __getitem__(self, val):
+        """
+        This is used when we want to index items.
+        """
+        return self.__item[val]
+
+    def __next__(self):
+        """
+        We want to make sure all iterators are still valid.
+        """
+        return self.__item.__next__()
+
+
+    def __iter__(self):
+        """
+        We want to make sure our wrapped items are still iterable.
+        """
+        return self.__item.__iter__()

src/hdmf/data_utils.py

@@ -1114,5 +1114,8 @@ def valid(self):
         return self.data is not None
 
 
+
+
+
 class InvalidDataIOError(Exception):
     pass

src/hdmf/utils.py

@@ -97,6 +97,8 @@
         elif argtype == 'bool':
             return __is_bool(value)
         return argtype in [cls.__name__ for cls in value.__class__.__mro__]
+    # elif isinstance(value, TermSetWrapper):
+        # pass
     elif isinstance(argtype, type):
         if argtype is int:
             return __is_int(value)
@@ -214,7 +216,6 @@
     future_warnings = list()
     argsi = 0
     extras = dict()  # has to be initialized to empty here, to avoid spurious errors reported upon early raises
-
     try:
         # check for duplicates in docval
         names = [x['name'] for x in validator]
@@ -273,6 +274,12 @@
                 type_errors.append("missing argument '%s'" % argname)
             else:
                 if enforce_type:
+                    from .container import TermSetWrapper # circular import fix
+                    termset = False
+                    if isinstance(argval, TermSetWrapper):
+                        # kwargs is the dict that stores the object names and the values
+                        # we can use this to unwrap the dataset/attribute to use the "item" for docval to validate the type.
+                        argval = kwargs[argname].item
                     if not __type_okay(argval, arg['type']):
                         if argval is None:
                             fmt_val = (argname, __format_type(arg['type']))