From 8fe396a61052df577551c07b10fb92b02314472e Mon Sep 17 00:00:00 2001
From: mavaylon1 <mavaylon1@berkeley.edu>
Date: Wed, 15 Nov 2023 08:18:04 -0800
Subject: [PATCH] termset schema option

---
 src/pynwb/core.py                      |  6 ++++++
 src/pynwb/file.py                      |  4 ++++
 src/pynwb/io/core.py                   | 10 +++++++++-
 src/pynwb/io/file.py                   |  8 +++++++-
 tests/integration/hdf5/test_nwbfile.py |  1 -
 5 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/pynwb/core.py b/src/pynwb/core.py
index b54f3e147..33a99fb8b 100644
--- a/src/pynwb/core.py
+++ b/src/pynwb/core.py
@@ -54,6 +54,12 @@ class NWBContainer(NWBMixin, Container):
 
     __nwbfields__ = tuple()
 
+@register_class('NWBTermSetContainer', CORE_NAMESPACE)
+class NWBTermSetContainer(NWBContainer):
+    def __init__(self, **kwargs):
+        termset_path = popargs('termset_path', kwargs)
+        super().__init__(**kwargs)
+        self.termset_path = termset_path
 
 @register_class('NWBDataInterface', CORE_NAMESPACE)
 class NWBDataInterface(NWBContainer):
diff --git a/src/pynwb/file.py b/src/pynwb/file.py
index b473e571a..e9146b459 100644
--- a/src/pynwb/file.py
+++ b/src/pynwb/file.py
@@ -107,6 +107,7 @@ class Subject(NWBContainer):
         {'name': 'date_of_birth', 'type': datetime, 'default': None,
          'doc': 'The datetime of the date of birth. May be supplied instead of age.'},
         {'name': 'strain', 'type': str, 'doc': 'The strain of the subject, e.g., "C57BL/6J"', 'default': None},
+        {'name': 'termset_path', 'type':str, 'doc': 'The path to the the TermSet', 'default': '...'}
     )
     def __init__(self, **kwargs):
         keys_to_set = (
@@ -121,9 +122,12 @@ def __init__(self, **kwargs):
             "date_of_birth",
             "strain",
         )
+        termset_path = popargs('termset_path', kwargs)
         args_to_set = popargs_to_dict(keys_to_set, kwargs)
         super().__init__(name="subject", **kwargs)
 
+        self.termset_path = termset_path
+
         # NOTE when the Subject I/O mapper (see pynwb.io.file.py) reads an age__reference value of None from an
         # NWB 2.0-2.5 file, it sets the value to "unspecified" so that when Subject.__init__ is called, the incoming
         # age__reference value is NOT replaced by the default value ("birth") specified in the docval.
diff --git a/src/pynwb/io/core.py b/src/pynwb/io/core.py
index be9b234d8..7f9614d1a 100644
--- a/src/pynwb/io/core.py
+++ b/src/pynwb/io/core.py
@@ -7,7 +7,7 @@
 from .. import register_map
 
 from pynwb.file import NWBFile
-from pynwb.core import NWBData, NWBContainer, ScratchData
+from pynwb.core import NWBData, NWBContainer, ScratchData, NWBTermSetContainer
 from pynwb.misc import Units
 
 
@@ -26,6 +26,14 @@ def get_nwb_file(container):
 class NWBContainerMapper(NWBBaseTypeMapper):
     pass
 
+@register_map(NWBTermSetContainer)
+class NWBTermSetContainerMapper(ObjectMapper):
+
+    def __init__(self, spec):
+        super().__init__(spec)
+        termset_path_spec = spec.get_attribute('termset_path')
+        self.map_spec('termset_path', termset_path_spec)
+
 
 @register_map(NWBData)
 class NWBDataMap(NWBBaseTypeMapper):
diff --git a/src/pynwb/io/file.py b/src/pynwb/io/file.py
index ccbfb8e47..0edd664fd 100644
--- a/src/pynwb/io/file.py
+++ b/src/pynwb/io/file.py
@@ -5,6 +5,7 @@
 from .. import register_map
 from ..file import NWBFile, Subject
 from ..core import ScratchData
+from .core import NWBTermSetContainerMapper
 from .utils import get_nwb_version
 
 
@@ -210,7 +211,12 @@ def publication_obj_attr(self, container, manager):
 
 
 @register_map(Subject)
-class SubjectMap(ObjectMapper):
+class SubjectMap(NWBTermSetContainerMapper):
+
+    # def __init__(self, spec):
+    #     super().__init__(spec)
+    #     termset_path_spec = self.spec.get_attribute('termset_path')
+    #     self.map_spec('termset_path', termset_path_spec)
 
     @ObjectMapper.constructor_arg('date_of_birth')
     def dateconversion(self, builder, manager):
diff --git a/tests/integration/hdf5/test_nwbfile.py b/tests/integration/hdf5/test_nwbfile.py
index e164ec649..3de9d1b95 100644
--- a/tests/integration/hdf5/test_nwbfile.py
+++ b/tests/integration/hdf5/test_nwbfile.py
@@ -499,7 +499,6 @@ def addContainer(self, nwbfile):
             group=self.group,
             group_name='tetrode1'
         )
-
         self.container = nwbfile.electrodes  # override self.container which has the placeholder
 
     def getContainer(self, nwbfile):