update Attributes, adding StoreV3 support

Author: grlee77

zarr/attrs.py

@@ -1,6 +1,6 @@
 from collections.abc import MutableMapping
 
-from zarr._storage.store import Store
+from zarr._storage.store import Store, StoreV3
 from zarr.util import json_dumps
 
 
@@ -26,7 +26,15 @@ class Attributes(MutableMapping):
 
     def __init__(self, store, key='.zattrs', read_only=False, cache=True,
                  synchronizer=None):
-        self.store = Store._ensure_store(store)
+
+        self._version = getattr(store, '_store_version', 2)
+        assert key
+
+        if self._version == 3 and '.z' in key:
+            raise ValueError('invalid v3 key')
+
+        _Store = Store if self._version == 2 else StoreV3
+        self.store = _Store._ensure_store(store)
         self.key = key
         self.read_only = read_only
         self.cache = cache
@@ -38,6 +46,8 @@ def _get_nosync(self):
             data = self.store[self.key]
         except KeyError:
             d = dict()
+            if self._version > 2:
+                d['attributes'] = {}
         else:
             d = self.store._metadata_class.parse_metadata(data)
         return d
@@ -47,14 +57,19 @@ def asdict(self):
         if self.cache and self._cached_asdict is not None:
             return self._cached_asdict
         d = self._get_nosync()
+        if self._version == 3:
+            d = d['attributes']
         if self.cache:
             self._cached_asdict = d
         return d
 
     def refresh(self):
         """Refresh cached attributes from the store."""
         if self.cache:
-            self._cached_asdict = self._get_nosync()
+            if self._version == 3:
+                self._cached_asdict = self._get_nosync()['attributes']
+            else:
+                self._cached_asdict = self._get_nosync()
 
     def __contains__(self, x):
         return x in self.asdict()
@@ -84,7 +99,10 @@ def _setitem_nosync(self, item, value):
         d = self._get_nosync()
 
         # set key value
-        d[item] = value
+        if self._version == 2:
+            d[item] = value
+        else:
+            d['attributes'][item] = value
 
         # _put modified data
         self._put_nosync(d)
@@ -98,20 +116,45 @@ def _delitem_nosync(self, key):
         d = self._get_nosync()
 
         # delete key value
-        del d[key]
+        if self._version == 2:
+            del d[key]
+        else:
+            del d['attributes'][key]
 
         # _put modified data
         self._put_nosync(d)
 
     def put(self, d):
         """Overwrite all attributes with the key/value pairs in the provided dictionary
         `d` in a single operation."""
-        self._write_op(self._put_nosync, d)
+        if self._version == 2:
+            self._write_op(self._put_nosync, d)
+        else:
+            self._write_op(self._put_nosync, dict(attributes=d))
 
     def _put_nosync(self, d):
-        self.store[self.key] = json_dumps(d)
-        if self.cache:
-            self._cached_asdict = d
+        if self._version == 2:
+            self.store[self.key] = json_dumps(d)
+            if self.cache:
+                self._cached_asdict = d
+        else:
+            if self.key in self.store:
+                # Cannot write the attributes directly to JSON, but have to
+                # store it within the pre-existing attributes key of the v3
+                # metadata.
+
+                # Note: this changes the store.counter result in test_caching_on!
+
+                meta = self.store._metadata_class.parse_metadata(self.store[self.key])
+                if 'attributes' in meta and 'filters' in meta['attributes']:
+                    # need to preserve any existing "filters" attribute
+                    d['attributes']['filters'] = meta['attributes']['filters']
+                meta['attributes'] = d['attributes']
+            else:
+                meta = d
+            self.store[self.key] = json_dumps(meta)
+            if self.cache:
+                self._cached_asdict = d['attributes']
 
     # noinspection PyMethodOverriding
     def update(self, *args, **kwargs):
@@ -124,7 +167,12 @@ def _update_nosync(self, *args, **kwargs):
         d = self._get_nosync()
 
         # update
-        d.update(*args, **kwargs)
+        if self._version == 2:
+            d.update(*args, **kwargs)
+        else:
+            if 'attributes' not in d:
+                d['attributes'] = {}
+            d['attributes'].update(*args, **kwargs)
 
         # _put modified data
         self._put_nosync(d)

zarr/tests/test_attrs.py

@@ -3,22 +3,30 @@
 import pytest
 
 from zarr.attrs import Attributes
-from zarr.tests.util import CountingDict
-from zarr.storage import KVStore
+from zarr.storage import KVStore, KVStoreV3
+from zarr.tests.util import CountingDict, CountingDictV3
+
+
+@pytest.fixture(params=[2, 3])
+def zarr_version(request):
+    return request.param
+
+
+def _init_store(version):
+    """Use a plain dict() for v2, but KVStoreV3 otherwise."""
+    if version == 2:
+        return dict()
+    return KVStoreV3(dict())
 
 
 class TestAttributes():
 
     def init_attributes(self, store, read_only=False, cache=True):
         return Attributes(store, key='attrs', read_only=read_only, cache=cache)
 
-    @pytest.mark.parametrize('store_from_dict', [False, True])
-    def test_storage(self, store_from_dict):
+    def test_storage(self, zarr_version):
 
-        if store_from_dict:
-            store = dict()
-        else:
-            store = KVStore(dict())
+        store = _init_store(zarr_version)
         a = Attributes(store=store, key='attrs')
         assert isinstance(a.store, KVStore)
         assert 'foo' not in a
@@ -30,11 +38,14 @@ def test_storage(self, store_from_dict):
         assert 'attrs' in store
         assert isinstance(store['attrs'], bytes)
         d = json.loads(str(store['attrs'], 'ascii'))
+        if zarr_version == 3:
+            d = d['attributes']
         assert dict(foo='bar', baz=42) == d
 
-    def test_get_set_del_contains(self):
+    def test_get_set_del_contains(self, zarr_version):
 
-        a = self.init_attributes(dict())
+        store = _init_store(zarr_version)
+        a = self.init_attributes(store)
         assert 'foo' not in a
         a['foo'] = 'bar'
         a['baz'] = 42
@@ -48,9 +59,10 @@ def test_get_set_del_contains(self):
             # noinspection PyStatementEffect
             a['foo']
 
-    def test_update_put(self):
+    def test_update_put(self, zarr_version):
 
-        a = self.init_attributes(dict())
+        store = _init_store(zarr_version)
+        a = self.init_attributes(store)
         assert 'foo' not in a
         assert 'bar' not in a
         assert 'baz' not in a
@@ -65,9 +77,10 @@ def test_update_put(self):
         assert a['bar'] == 84
         assert 'baz' not in a
 
-    def test_iterators(self):
+    def test_iterators(self, zarr_version):
 
-        a = self.init_attributes(dict())
+        store = _init_store(zarr_version)
+        a = self.init_attributes(store)
         assert 0 == len(a)
         assert set() == set(a)
         assert set() == set(a.keys())
@@ -83,10 +96,13 @@ def test_iterators(self):
         assert {'bar', 42} == set(a.values())
         assert {('foo', 'bar'), ('baz', 42)} == set(a.items())
 
-    def test_read_only(self):
-        store = dict()
+    def test_read_only(self, zarr_version):
+        store = _init_store(zarr_version)
         a = self.init_attributes(store, read_only=True)
-        store['attrs'] = json.dumps(dict(foo='bar', baz=42)).encode('ascii')
+        if zarr_version == 2:
+            store['attrs'] = json.dumps(dict(foo='bar', baz=42)).encode('ascii')
+        else:
+            store['attrs'] = json.dumps(dict(attributes=dict(foo='bar', baz=42))).encode('ascii')
         assert a['foo'] == 'bar'
         assert a['baz'] == 42
         with pytest.raises(PermissionError):
@@ -96,8 +112,9 @@ def test_read_only(self):
         with pytest.raises(PermissionError):
             a.update(foo='quux')
 
-    def test_key_completions(self):
-        a = self.init_attributes(dict())
+    def test_key_completions(self, zarr_version):
+        store = _init_store(zarr_version)
+        a = self.init_attributes(store)
         d = a._ipython_key_completions_()
         assert 'foo' not in d
         assert '123' not in d
@@ -112,14 +129,17 @@ def test_key_completions(self):
         assert 'asdf;' in d
         assert 'baz' not in d
 
-    def test_caching_on(self):
+    def test_caching_on(self, zarr_version):
         # caching is turned on by default
 
         # setup store
-        store = CountingDict()
+        store = CountingDict() if zarr_version == 2 else CountingDictV3()
         assert 0 == store.counter['__getitem__', 'attrs']
         assert 0 == store.counter['__setitem__', 'attrs']
-        store['attrs'] = json.dumps(dict(foo='xxx', bar=42)).encode('ascii')
+        if zarr_version == 2:
+            store['attrs'] = json.dumps(dict(foo='xxx', bar=42)).encode('ascii')
+        else:
+            store['attrs'] = json.dumps(dict(attributes=dict(foo='xxx', bar=42))).encode('ascii')
         assert 0 == store.counter['__getitem__', 'attrs']
         assert 1 == store.counter['__setitem__', 'attrs']
 
@@ -136,54 +156,65 @@ def test_caching_on(self):
 
         # test __setitem__ updates the cache
         a['foo'] = 'yyy'
-        assert 2 == store.counter['__getitem__', 'attrs']
+        get_cnt = 2 if zarr_version == 2 else 3
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 2 == store.counter['__setitem__', 'attrs']
         assert a['foo'] == 'yyy'
-        assert 2 == store.counter['__getitem__', 'attrs']
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 2 == store.counter['__setitem__', 'attrs']
 
         # test update() updates the cache
         a.update(foo='zzz', bar=84)
-        assert 3 == store.counter['__getitem__', 'attrs']
+        get_cnt = 3 if zarr_version == 2 else 5
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 3 == store.counter['__setitem__', 'attrs']
         assert a['foo'] == 'zzz'
         assert a['bar'] == 84
-        assert 3 == store.counter['__getitem__', 'attrs']
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 3 == store.counter['__setitem__', 'attrs']
 
         # test __contains__ uses the cache
         assert 'foo' in a
-        assert 3 == store.counter['__getitem__', 'attrs']
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 3 == store.counter['__setitem__', 'attrs']
         assert 'spam' not in a
-        assert 3 == store.counter['__getitem__', 'attrs']
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 3 == store.counter['__setitem__', 'attrs']
 
         # test __delitem__ updates the cache
         del a['bar']
-        assert 4 == store.counter['__getitem__', 'attrs']
+        get_cnt = 4 if zarr_version == 2 else 7
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 4 == store.counter['__setitem__', 'attrs']
         assert 'bar' not in a
-        assert 4 == store.counter['__getitem__', 'attrs']
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 4 == store.counter['__setitem__', 'attrs']
 
         # test refresh()
-        store['attrs'] = json.dumps(dict(foo='xxx', bar=42)).encode('ascii')
-        assert 4 == store.counter['__getitem__', 'attrs']
+        if zarr_version == 2:
+            store['attrs'] = json.dumps(dict(foo='xxx', bar=42)).encode('ascii')
+        else:
+            store['attrs'] = json.dumps(dict(attributes=dict(foo='xxx', bar=42))).encode('ascii')
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         a.refresh()
-        assert 5 == store.counter['__getitem__', 'attrs']
+        get_cnt = 5 if zarr_version == 2 else 8
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert a['foo'] == 'xxx'
-        assert 5 == store.counter['__getitem__', 'attrs']
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert a['bar'] == 42
-        assert 5 == store.counter['__getitem__', 'attrs']
+        assert get_cnt == store.counter['__getitem__', 'attrs']
 
-    def test_caching_off(self):
+    def test_caching_off(self, zarr_version):
 
         # setup store
-        store = CountingDict()
+        store = CountingDict() if zarr_version == 2 else CountingDictV3()
         assert 0 == store.counter['__getitem__', 'attrs']
         assert 0 == store.counter['__setitem__', 'attrs']
-        store['attrs'] = json.dumps(dict(foo='xxx', bar=42)).encode('ascii')
+
+        if zarr_version == 2:
+            store['attrs'] = json.dumps(dict(foo='xxx', bar=42)).encode('ascii')
+        else:
+            store['attrs'] = json.dumps(dict(attributes=dict(foo='xxx', bar=42))).encode('ascii')
         assert 0 == store.counter['__getitem__', 'attrs']
         assert 1 == store.counter['__setitem__', 'attrs']
 
@@ -200,25 +231,31 @@ def test_caching_off(self):
 
         # test __setitem__
         a['foo'] = 'yyy'
-        assert 4 == store.counter['__getitem__', 'attrs']
+        get_cnt = 4 if zarr_version == 2 else 5
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 2 == store.counter['__setitem__', 'attrs']
         assert a['foo'] == 'yyy'
-        assert 5 == store.counter['__getitem__', 'attrs']
+        get_cnt = 5 if zarr_version == 2 else 6
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 2 == store.counter['__setitem__', 'attrs']
 
         # test update()
         a.update(foo='zzz', bar=84)
-        assert 6 == store.counter['__getitem__', 'attrs']
+        get_cnt = 6 if zarr_version == 2 else 8
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 3 == store.counter['__setitem__', 'attrs']
         assert a['foo'] == 'zzz'
         assert a['bar'] == 84
-        assert 8 == store.counter['__getitem__', 'attrs']
+        get_cnt = 8 if zarr_version == 2 else 10
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 3 == store.counter['__setitem__', 'attrs']
 
         # test __contains__
         assert 'foo' in a
-        assert 9 == store.counter['__getitem__', 'attrs']
+        get_cnt = 9 if zarr_version == 2 else 11
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 3 == store.counter['__setitem__', 'attrs']
         assert 'spam' not in a
-        assert 10 == store.counter['__getitem__', 'attrs']
+        get_cnt = 10 if zarr_version == 2 else 12
+        assert get_cnt == store.counter['__getitem__', 'attrs']
         assert 3 == store.counter['__setitem__', 'attrs']