[WIP] Add/zarrio

src/pynwb/__init__.py

@@ -13,6 +13,7 @@
 from hdmf.utils import docval, getargs, popargs, call_docval_func  # noqa: E402
 from hdmf.backends.io import HDMFIO  # noqa: E402
 from hdmf.backends.hdf5 import HDF5IO as _HDF5IO  # noqa: E402
+from hdmf.backends.zarr.zarr_tools import ZarrIO as _ZarrIO
 from hdmf.validate import ValidatorMap  # noqa: E402
 from hdmf.build import BuildManager  # noqa: E402
 
@@ -231,6 +232,44 @@ def __init__(self, **kwargs):
                 manager = get_manager()
         super(NWBHDF5IO, self).__init__(path, manager=manager, mode=mode, file=file_obj, comm=comm)
 
+class NWBZarrIO(_ZarrIO):
+
+    @docval({'name': 'path', 'type': str, 'doc': 'the path to the Zarr file'},
+            {'name': 'mode', 'type': str,
+             'doc': 'the mode to open the Zarr file with, one of ("w", "r", "r+", "a", "w-")'},
+            {'name': 'load_namespaces', 'type': bool,
+             'doc': 'whether or not to load cached namespaces from given path - not applicable in write mode',
+             'default': False},
+            {'name': 'manager', 'type': BuildManager, 'doc': 'the BuildManager to use for I/O', 'default': None},
+            {'name': 'extensions', 'type': (str, TypeMap, list),
+              'doc': 'a path to a namespace, a TypeMap, or a list consisting paths \
+              to namespaces and TypeMaps', 'default': None},
+            {'name': 'comm', 'type': 'Intracomm',
+             'doc': 'the MPI communicator to use for parallel I/O', 'default': None})
+    def __init__(self, **kwargs):
+        path, mode, manager, extensions, load_namespaces, comm =\
+            popargs('path', 'mode', 'manager', 'extensions', 'load_namespaces', 'comm', kwargs)
+        if load_namespaces:
+            if manager is not None:
+                warn("loading namespaces from file - ignoring 'manager'")
+            if extensions is not None:
+                warn("loading namespaces from file - ignoring 'extensions' argument")
+            # namespaces are not loaded when creating an NWBHDF5IO object in write mode
+            if 'w' in mode or mode == 'x':
+                raise ValueError("cannot load namespaces from file when writing to it")
+
+            tm = get_type_map()
+            super(NWBZarrIO, self).load_namespaces(tm, path)
+            manager = BuildManager(tm)
+        else:
+            if manager is not None and extensions is not None:
+                raise ValueError("'manager' and 'extensions' cannot be specified together")
+            elif extensions is not None:
+                manager = get_manager(extensions=extensions)
+            elif manager is None:
+                manager = get_manager()
+        super(NWBZarrIO, self).__init__(path, manager=manager, mode=mode, comm=comm)
+
 
 from . import io as __io  # noqa: F401,E402
 from .core import NWBContainer, NWBData  # noqa: F401,E402

tests/integration/ui_write/base.py

@@ -4,10 +4,14 @@
 import os
 import numpy as np
 import h5py
+import shutil
 import numpy.testing as npt
 
-from pynwb import NWBContainer, get_manager, NWBFile, NWBData, NWBHDF5IO, validate as pynwb_validate
+from pynwb import NWBContainer, get_manager, NWBFile, NWBData, NWBHDF5IO, NWBZarrIO, validate as pynwb_validate
 from hdmf.backends.hdf5 import HDF5IO
+from hdmf.backends.zarr import ZarrIO
+
+from zarr.core import Array as ZarrArray
 
 CORE_NAMESPACE = 'core'
 
@@ -91,6 +95,8 @@ def assertContainerEqual(self, container1, container2):           # noqa: C901
                 f2 = getattr(container2, nwbfield)
                 if isinstance(f1, h5py.Dataset):
                     f1 = f1[()]
+                if isinstance(f1, ZarrArray):
+                    f1 = f1[:]
                 if isinstance(f1, (tuple, list, np.ndarray)):
                     if len(f1) > 0:
                         if isinstance(f1[0], NWBContainer):
@@ -147,7 +153,11 @@ def setUp(self):
         self.start_time = datetime(1971, 1, 1, 12, tzinfo=tzutc())
         self.create_date = datetime(2018, 4, 15, 12, tzinfo=tzlocal())
         self.container_type = self.container.__class__.__name__
-        self.filename = 'test_%s.nwb' % self.container_type
+        test_case_name = str(self.id()).split(".")[-1]
+        if test_case_name == "test_zarr_roundtrip":
+            self.filename = 'test_zarrio_%s' % self.container_type
+        else:
+            self.filename = 'test_%s.nwb' % self.container_type
         self.writer = None
         self.reader = None
 
@@ -157,7 +167,10 @@ def tearDown(self):
         if self.reader is not None:
             self.reader.close()
         if os.path.exists(self.filename) and os.getenv("CLEAN_NWB", '1') not in ('0', 'false', 'FALSE', 'False'):
-            os.remove(self.filename)
+            if os.path.isfile(self.filename):
+                os.remove(self.filename)
+            elif os.path.isdir(self.filename):
+                shutil.rmtree(self.filename)
 
     def roundtripContainer(self, cache_spec=False):
         description = 'a file to test writing and reading a %s' % self.container_type
@@ -196,6 +209,45 @@ def validate(self):
                     for err in errors:
                         raise Exception(err)
 
+    def roundtripContainerZarrIO(self, cache_spec=False):
+        description = 'a file to test writing and reading a %s' % self.container_type
+        identifier = 'TEST_%s' % self.container_type
+        nwbfile = NWBFile(description, identifier, self.start_time, file_create_date=self.create_date)
+        self.addContainer(nwbfile)
+
+        self.writer = ZarrIO(self.filename, manager=get_manager(), mode='w')
+        self.writer.write(nwbfile, cache_spec=cache_spec)
+        self.writer.close()
+        self.reader = ZarrIO(self.filename, manager=get_manager(), mode='r')
+        self.read_nwbfile = self.reader.read()
+
+        try:
+            tmp = self.getContainer(self.read_nwbfile)
+            return tmp
+        except Exception as e:
+            self.reader.close()
+            self.reader = None
+            raise e
+
+    def test_zarr_roundtrip(self):
+        self.read_container = self.roundtripContainerZarrIO()
+        # make sure we get a completely new object
+        str(self.container)  # added as a test to make sure printing works
+        self.assertNotEqual(id(self.container), id(self.read_container))
+        self.assertContainerEqual(self.read_container, self.container)
+        self.validate_zarr()
+
+    def validate_zarr(self):
+        # validate created file
+        if os.path.exists(self.filename):
+            with NWBZarrIO(self.filename, mode='r') as io:
+                # TODO need to update the validator to support Zarr. For now well just read the file instead
+                #errors = pynwb_validate(io)
+                #if errors:
+                #    for err in errors:
+                #        raise Exception(err)
+                temp = io.read()
+
     def addContainer(self, nwbfile):
         ''' Should take an NWBFile object and add the container to it '''
         raise unittest.SkipTest('Cannot run test unless addContainer is implemented')

tests/integration/ui_write/test_modular_storage.py

@@ -22,6 +22,10 @@ def remove_file(self, path):
             os.remove(path)
 
     def setUp(self):
+        test_case_name = str(self.id()).split(".")[-1]
+        if test_case_name == "test_zarr_roundtrip":
+            self.skipTest("Modular storage testing does not apply to ZarrIO")
+
         self.__manager = get_manager()
         self.start_time = datetime(1971, 1, 1, 12, tzinfo=tzutc())