Type dimension separator (#1620)

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

Author: dstansby

zarr/_storage/absstore.py

@@ -1,9 +1,12 @@
 """This module contains storage classes related to Azure Blob Storage (ABS)"""
 
+from typing import Optional
 import warnings
+
 from numcodecs.compat import ensure_bytes
 from zarr.util import normalize_storage_path
 from zarr._storage.store import _get_metadata_suffix, data_root, meta_root, Store, StoreV3
+from zarr.types import DIMENSION_SEPARATOR
 
 __doctest_requires__ = {
     ("ABSStore", "ABSStore.*"): ["azure.storage.blob"],
@@ -67,7 +70,7 @@ def __init__(
         account_name=None,
         account_key=None,
         blob_service_kwargs=None,
-        dimension_separator=None,
+        dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
         client=None,
     ):
         self._dimension_separator = dimension_separator

zarr/_storage/v3.py

@@ -3,13 +3,14 @@
 from collections import OrderedDict
 from collections.abc import MutableMapping
 from threading import Lock
-from typing import Union, Dict, Any
+from typing import Union, Dict, Any, Optional
 
 from zarr.errors import (
     MetadataError,
     ReadOnlyError,
 )
 from zarr.util import buffer_size, json_loads, normalize_storage_path
+from zarr.types import DIMENSION_SEPARATOR
 
 from zarr._storage.absstore import ABSStoreV3  # noqa: F401
 from zarr._storage.store import (  # noqa: F401
@@ -224,7 +225,9 @@ def get_partial_values(self, key_ranges):
 
 
 class MemoryStoreV3(MemoryStore, StoreV3):
-    def __init__(self, root=None, cls=dict, dimension_separator=None):
+    def __init__(
+        self, root=None, cls=dict, dimension_separator: Optional[DIMENSION_SEPARATOR] = None
+    ):
         if root is None:
             self.root = cls()
         else:

zarr/_storage/v3_storage_transformers.py

@@ -8,6 +8,7 @@
 
 from zarr._storage.store import StorageTransformer, StoreV3, _rmdir_from_keys_v3
 from zarr.util import normalize_storage_path
+from zarr.types import DIMENSION_SEPARATOR
 
 
 MAX_UINT_64 = 2**64 - 1
@@ -118,7 +119,7 @@ def _copy_for_array(self, array, inner_store):
         return transformer_copy
 
     @property
-    def dimension_separator(self) -> str:
+    def dimension_separator(self) -> DIMENSION_SEPARATOR:
         assert (
             self._dimension_separator is not None
         ), "dimension_separator is not initialized, first get a copy via _copy_for_array."

zarr/creation.py

@@ -470,7 +470,7 @@ def open_array(
     write_empty_chunks=True,
     *,
     zarr_version=None,
-    dimension_separator=None,
+    dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
     meta_array=None,
     **kwargs,
 ):

zarr/storage.py

@@ -41,7 +41,8 @@
 from numcodecs.compat import ensure_bytes, ensure_text, ensure_contiguous_ndarray_like
 from numcodecs.registry import codec_registry
 from zarr.context import Context
-from zarr.types import PathLike as Path
+from zarr.types import PathLike as Path, DIMENSION_SEPARATOR
+from zarr.util import NoLock
 
 from zarr.errors import (
     MetadataError,
@@ -327,7 +328,7 @@ def init_array(
     chunk_store: Optional[StoreLike] = None,
     filters=None,
     object_codec=None,
-    dimension_separator=None,
+    dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
     storage_transformers=(),
 ):
     """Initialize an array store with the given configuration. Note that this is a low-level
@@ -481,7 +482,7 @@ def _init_array_metadata(
     chunk_store: Optional[StoreLike] = None,
     filters=None,
     object_codec=None,
-    dimension_separator=None,
+    dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
     storage_transformers=(),
 ):
     store_version = getattr(store, "_store_version", 2)
@@ -1054,7 +1055,9 @@ class DirectoryStore(Store):
 
     """
 
-    def __init__(self, path, normalize_keys=False, dimension_separator=None):
+    def __init__(
+        self, path, normalize_keys=False, dimension_separator: Optional[DIMENSION_SEPARATOR] = None
+    ):
         # guard conditions
         path = os.path.abspath(path)
         if os.path.exists(path) and not os.path.isdir(path):
@@ -1349,7 +1352,7 @@ def __init__(
         key_separator=None,
         mode="w",
         exceptions=(KeyError, PermissionError, IOError),
-        dimension_separator=None,
+        dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
         fs=None,
         check=False,
         create=False,
@@ -1568,7 +1571,12 @@ class TempStore(DirectoryStore):
 
     # noinspection PyShadowingBuiltins
     def __init__(
-        self, suffix="", prefix="zarr", dir=None, normalize_keys=False, dimension_separator=None
+        self,
+        suffix="",
+        prefix="zarr",
+        dir=None,
+        normalize_keys=False,
+        dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
     ):
         path = tempfile.mkdtemp(suffix=suffix, prefix=prefix, dir=dir)
         atexit.register(atexit_rmtree, path)
@@ -1652,7 +1660,9 @@ class NestedDirectoryStore(DirectoryStore):
 
     """
 
-    def __init__(self, path, normalize_keys=False, dimension_separator="/"):
+    def __init__(
+        self, path, normalize_keys=False, dimension_separator: Optional[DIMENSION_SEPARATOR] = "/"
+    ):
         super().__init__(path, normalize_keys=normalize_keys)
         if dimension_separator is None:
             dimension_separator = "/"
@@ -1765,7 +1775,7 @@ def __init__(
         compression=zipfile.ZIP_STORED,
         allowZip64=True,
         mode="a",
-        dimension_separator=None,
+        dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
     ):
         # store properties
         path = os.path.abspath(path)
@@ -2058,7 +2068,7 @@ def __init__(
         mode=0o666,
         open=None,
         write_lock=True,
-        dimension_separator=None,
+        dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
         **open_kwargs,
     ):
         if open is None:
@@ -2073,6 +2083,7 @@ def __init__(
         self.mode = mode
         self.open = open
         self.write_lock = write_lock
+        self.write_mutex: Union[Lock, NoLock]
         if write_lock:
             # This may not be required as some dbm implementations manage their own
             # locks, but err on the side of caution.
@@ -2229,7 +2240,13 @@ class LMDBStore(Store):
 
     """
 
-    def __init__(self, path, buffers=True, dimension_separator=None, **kwargs):
+    def __init__(
+        self,
+        path,
+        buffers=True,
+        dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
+        **kwargs,
+    ):
         import lmdb
 
         # set default memory map size to something larger than the lmdb default, which is
@@ -2580,7 +2597,7 @@ class SQLiteStore(Store):
         >>> store.close()  # don't forget to call this when you're done
     """
 
-    def __init__(self, path, dimension_separator=None, **kwargs):
+    def __init__(self, path, dimension_separator: Optional[DIMENSION_SEPARATOR] = None, **kwargs):
         import sqlite3
 
         self._dimension_separator = dimension_separator
@@ -2776,7 +2793,7 @@ def __init__(
         self,
         database="mongodb_zarr",
         collection="zarr_collection",
-        dimension_separator=None,
+        dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
         **kwargs,
     ):
         import pymongo
@@ -2851,7 +2868,9 @@ class RedisStore(Store):
 
     """
 
-    def __init__(self, prefix="zarr", dimension_separator=None, **kwargs):
+    def __init__(
+        self, prefix="zarr", dimension_separator: Optional[DIMENSION_SEPARATOR] = None, **kwargs
+    ):
         import redis
 
         self._prefix = prefix