Make Index.setdefault atomic using Cache.add

Author: grantjenks

diskcache/persistent.py

@@ -728,6 +728,7 @@ def __getitem__(self, key):
  KeyError: 'c'
 
  :param key: key for item
+ :return: value for item in index with given key
  :raises KeyError: if key is not found
 
  """
@@ -799,6 +800,39 @@ def __delitem__(self, key):
 return
 
 
+defsetdefault(self, key, default=None):
+"""Set and get value for `key` in index using `default`.
+
+ If `key` is not in index then set corresponding value to `default`. If
+ `key` is in index then ignore `default` and return existing value.
+
+ >>> index = Index('/tmp/diskcache/index')
+ >>> index.clear()
+ >>> index.setdefault('a', 0)
+ 0
+ >>> index.setdefault('a', 1)
+ 0
+
+ :param key: key for item
+ :param default: value if key is missing (default None)
+ :return: value for item in index with given key
+
+ """
+_cache=self._cache
+
+whileTrue:
+try:
+returnself[key]
+exceptKeyError:
+whileTrue:
+try:
+_cache.add(key, default)
+exceptTimeout:
+continue
+else:
+break
+
+
 defpop(self, key, default=ENOVAL):
 """Remove corresponding item for `key` from index and return value.
 

diskcache/stampede.py

@@ -21,10 +21,13 @@ class StampedeBarrier(object):
 
  Example:
 
- >>> stampede_barrier = StampedeBarrier('/tmp/user_data', expire=3)
- >>> @stampede_barrier
+ ```python
+ stampede_barrier = StampedeBarrier('/tmp/user_data', expire=3)
+
+ @stampede_barrier
  def load_user_info(user_id):
  return database.lookup_user_info_by_id(user_id)
+ ```
 
  """
 # pylint: disable=too-few-public-methods

tests/test_index.py

@@ -170,6 +170,25 @@ def test_popitem_timeout(index):
 assertvalue== (0, 1)
 
 
+@setup_index
+deftest_setdefault(index):
+assertindex.setdefault('a', 0) ==0
+assertindex.setdefault('a', 1) ==0
+
+
+@setup_index
+deftest_setdefault_timeout(index):
+cache=mock.MagicMock()
+cache.__getitem__=mock.Mock()
+cache.__getitem__.side_effect= [KeyError, 0]
+cache.add=mock.Mock()
+cache.add.side_effect= [dc.Timeout, 0]
+
+withmock.patch.object(index, '_cache', cache):
+value=index.setdefault('a', 0)
+assertvalue==0
+
+
 @setup_index
 deftest_iter(index):
 letters='abcde'