Make DirectoryStore dimension_separator aware

zarr/storage.py

@@ -803,6 +803,10 @@ def __init__(self, path, normalize_keys=False, dimension_separator=None):
     def _normalize_key(self, key):
         return key.lower() if self.normalize_keys else key
 
+    def _optionally_nested(self, key):
+        return self._dimension_separator == "/" and \
+            _nested_map_ckey(key) or key
+
     def _fromfile(self, fn):
         """ Read data from a file
 
@@ -838,6 +842,7 @@ def _tofile(self, a, fn):
             f.write(a)
 
     def __getitem__(self, key):
+        key = self._optionally_nested(key)
         key = self._normalize_key(key)
         filepath = os.path.join(self.path, key)
         if os.path.isfile(filepath):
@@ -846,6 +851,7 @@ def __getitem__(self, key):
             raise KeyError(key)
 
     def __setitem__(self, key, value):
+        key = self._optionally_nested(key)
         key = self._normalize_key(key)
 
         # coerce to flat, contiguous array (ideally without copying)
@@ -887,6 +893,7 @@ def __setitem__(self, key, value):
                 os.remove(temp_path)
 
     def __delitem__(self, key):
+        key = self._optionally_nested(key)
         key = self._normalize_key(key)
         path = os.path.join(self.path, key)
         if os.path.isfile(path):
@@ -899,6 +906,7 @@ def __delitem__(self, key):
             raise KeyError(key)
 
     def __contains__(self, key):
+        key = self._optionally_nested(key)
         key = self._normalize_key(key)
         file_path = os.path.join(self.path, key)
         return os.path.isfile(file_path)
@@ -947,12 +955,37 @@ def dir_path(self, path=None):
         return dir_path
 
     def listdir(self, path=None):
+        return self._dimension_separator == "/" and \
+            self._nested_listdir(path) or self._flat_listdir(path)
+
+    def _flat_listdir(self, path=None):
         dir_path = self.dir_path(path)
         if os.path.isdir(dir_path):
             return sorted(os.listdir(dir_path))
         else:
             return []
 
+    def _nested_listdir(self, path=None):
+        children = self._flat_listdir(path=path)
+        if array_meta_key in children:
+            # special handling of directories containing an array to map nested chunk
+            # keys back to standard chunk keys
+            new_children = []
+            root_path = self.dir_path(path)
+            for entry in children:
+                entry_path = os.path.join(root_path, entry)
+                if _prog_number.match(entry) and os.path.isdir(entry_path):
+                    for dir_path, _, file_names in os.walk(entry_path):
+                        for file_name in file_names:
+                            file_path = os.path.join(dir_path, file_name)
+                            rel_path = file_path.split(root_path + os.path.sep)[1]
+                            new_children.append(rel_path.replace(os.path.sep, '.'))
+                else:
+                    new_children.append(entry)
+            return sorted(new_children)
+        else:
+            return children
+
     def rename(self, src_path, dst_path):
         store_src_path = normalize_storage_path(src_path)
         store_dst_path = normalize_storage_path(dst_path)
@@ -1314,49 +1347,12 @@ def __init__(self, path, normalize_keys=False, dimension_separator="/"):
                 "NestedDirectoryStore only supports '/' as dimension_separator")
         self._dimension_separator = dimension_separator
 
-    def __getitem__(self, key):
-        key = _nested_map_ckey(key)
-        return super().__getitem__(key)
-
-    def __setitem__(self, key, value):
-        key = _nested_map_ckey(key)
-        super().__setitem__(key, value)
-
-    def __delitem__(self, key):
-        key = _nested_map_ckey(key)
-        super().__delitem__(key)
-
-    def __contains__(self, key):
-        key = _nested_map_ckey(key)
-        return super().__contains__(key)
-
     def __eq__(self, other):
         return (
             isinstance(other, NestedDirectoryStore) and
             self.path == other.path
         )
 
-    def listdir(self, path=None):
-        children = super().listdir(path=path)
-        if array_meta_key in children:
-            # special handling of directories containing an array to map nested chunk
-            # keys back to standard chunk keys
-            new_children = []
-            root_path = self.dir_path(path)
-            for entry in children:
-                entry_path = os.path.join(root_path, entry)
-                if _prog_number.match(entry) and os.path.isdir(entry_path):
-                    for dir_path, _, file_names in os.walk(entry_path):
-                        for file_name in file_names:
-                            file_path = os.path.join(dir_path, file_name)
-                            rel_path = file_path.split(root_path + os.path.sep)[1]
-                            new_children.append(rel_path.replace(os.path.sep, '.'))
-                else:
-                    new_children.append(entry)
-            return sorted(new_children)
-        else:
-            return children
-
 
 # noinspection PyPep8Naming
 class ZipStore(MutableMapping):