add tests for keys caching

Author: alimanfoo

zarr/storage.py

@@ -1696,18 +1696,21 @@ def __init__(self, store, max_size):
         self._max_size = max_size
         self._current_size = 0
         self._keys_cache = None
+        self._contains_cache = None
         self._listdir_cache = dict()
         self._values_cache = OrderedDict()
         self._mutex = Lock()
         self.hits = self.misses = 0
 
     def __getstate__(self):
         return (self._store, self._max_size, self._current_size, self._keys_cache,
-                self._listdir_cache, self._values_cache, self.hits, self.misses)
+                self._contains_cache, self._listdir_cache, self._values_cache, self.hits,
+                self.misses)
 
     def __setstate__(self, state):
         (self._store, self._max_size, self._current_size, self._keys_cache,
-         self._listdir_cache, self._values_cache, self.hits, self.misses) = state
+         self._contains_cache, self._listdir_cache, self._values_cache, self.hits,
+         self.misses) = state
         self._mutex = Lock()
 
     def __len__(self):
@@ -1716,10 +1719,19 @@ def __len__(self):
     def __iter__(self):
         return self.keys()
 
+    def __contains__(self, key):
+        with self._mutex:
+            if self._contains_cache is None:
+                self._contains_cache = set(self._keys())
+            return key in self._contains_cache
+
     def keys(self):
         with self._mutex:
-            if self._keys_cache is None:
-                self._keys_cache = list(self._store.keys())
+            return self._keys()
+
+    def _keys(self):
+        if self._keys_cache is None:
+            self._keys_cache = list(self._store.keys())
         return iter(self._keys_cache)
 
     def listdir(self, path=None):
@@ -1768,6 +1780,7 @@ def clear_keys(self):
 
     def _clear_keys(self):
         self._keys_cache = None
+        self._contains_cache = None
         self._listdir_cache.clear()
 
     def _clear_value(self, key):

zarr/tests/test_storage.py

@@ -889,6 +889,12 @@ def test_cache_values_no_max_size(self):
         assert 2 == cache.hits
         assert 1 == cache.misses
 
+        # manually clear all cached values
+        cache.clear_values()
+        assert b'zzz' == cache['foo']
+        assert 2 == store.counter['__getitem__', 'foo']
+        assert 2 == store.counter['__setitem__', 'foo']
+
         # test __delitem__
         del cache['foo']
         with pytest.raises(KeyError):
@@ -902,7 +908,6 @@ def test_cache_values_no_max_size(self):
         assert 0 == store.counter['__getitem__', 'bar']
         assert 1 == store.counter['__setitem__', 'bar']
 
-    # TODO test max size
     def test_cache_values_with_max_size(self):
 
         # setup store
@@ -999,7 +1004,63 @@ def test_cache_values_with_max_size(self):
         assert 4 == cache.hits
         assert 2 == cache.misses
 
-    # TODO test key caching
+    def test_cache_keys(self):
+
+        # setup
+        store = CountingDict()
+        store['foo'] = b'xxx'
+        store['bar'] = b'yyy'
+        assert 0 == store.counter['__contains__', 'foo']
+        assert 0 == store.counter['__iter__']
+        assert 0 == store.counter['keys']
+        cache = LRUStoreCache(store, max_size=None)
+
+        # keys should be cached on first call
+        keys = sorted(cache.keys())
+        assert keys == ['bar', 'foo']
+        assert 1 == store.counter['keys']
+        # keys should now be cached
+        assert keys == sorted(cache.keys())
+        assert 1 == store.counter['keys']
+        assert 'foo' in cache
+        assert 0 == store.counter['__contains__', 'foo']
+        assert keys == sorted(cache)
+        assert 0 == store.counter['__iter__']
+        assert 1 == store.counter['keys']
+
+        # cache should be cleared if store is modified - crude but simple for now
+        cache['baz'] = b'zzz'
+        keys = sorted(cache.keys())
+        assert keys == ['bar', 'baz', 'foo']
+        assert 2 == store.counter['keys']
+        # keys should now be cached
+        assert keys == sorted(cache.keys())
+        assert 2 == store.counter['keys']
+
+        # manually clear keys
+        cache.clear_keys()
+        keys = sorted(cache.keys())
+        assert keys == ['bar', 'baz', 'foo']
+        assert 3 == store.counter['keys']
+        assert 0 == store.counter['__contains__', 'foo']
+        assert 0 == store.counter['__iter__']
+        cache.clear_keys()
+        keys = sorted(cache)
+        assert keys == ['bar', 'baz', 'foo']
+        assert 4 == store.counter['keys']
+        assert 0 == store.counter['__contains__', 'foo']
+        assert 0 == store.counter['__iter__']
+        cache.clear_keys()
+        assert 'foo' in cache
+        assert 5 == store.counter['keys']
+        assert 0 == store.counter['__contains__', 'foo']
+        assert 0 == store.counter['__iter__']
+
+        # check these would get counted if called directly
+        assert 'foo' in store
+        assert 1 == store.counter['__contains__', 'foo']
+        assert keys == sorted(store)
+        assert 1 == store.counter['__iter__']
 
 
 def test_getsize():

zarr/tests/util.py

@@ -10,12 +10,19 @@ def __init__(self):
         self.counter = collections.Counter()
 
     def __len__(self):
+        self.counter['__len__'] += 1
         return len(self.wrapped)
 
+    def keys(self):
+        self.counter['keys'] += 1
+        return self.wrapped.keys()
+
     def __iter__(self):
+        self.counter['__iter__'] += 1
         return iter(self.wrapped)
 
     def __contains__(self, item):
+        self.counter['__contains__', item] += 1
         return item in self.wrapped
 
     def __getitem__(self, item):