Switch to caching with pickle instead of shelve

.gitignore

@@ -6,4 +6,4 @@ dist/
 budou.egg-info/
 .eggs/
 *.db
-budou-cache.shelve
+budou-cache.pickle

README.md

@@ -78,7 +78,7 @@ API and make the processing faster. If you want to force refresh the cache,
 put `use_cache=False`.
 
 In a standard environment, Budou will create a cache file with
-[python shelve](https://docs.python.org/3/library/shelve.html) format.
+[python pickle](https://docs.python.org/3/library/pickle.html) format.
 
 In [Google App Engine Python Standard Environment](https://cloud.google.com/appengine/docs/standard/python/),
 Budou will use [memcache](https://cloud.google.com/appengine/docs/standard/python/memcache/)

budou/cachefactory.py

@@ -16,14 +16,15 @@
 from abc import ABCMeta, abstractmethod
 import hashlib
 import six
-import shelve
+import pickle
+import os
 
 
 def load_cache():
   try:
     from google.appengine.api import memcache
   except:
-    return ShelveCache()
+    return PickleCache()
   else:
     return AppEngineCache(memcache)
 
@@ -32,7 +33,7 @@ def load_cache():
 class BudouCache(object):
 
   CACHE_SALT = '2017-04-13'
-  DEFAULT_FILE_PATH = './budou-cache.shelve'
+  DEFAULT_FILE_PATH = './budou-cache.pickle'
 
   def __repr__(self):
     return '<%s>' % (self.__class__.__name__)
@@ -51,23 +52,33 @@ def _get_cache_key(self, source, language):
     return hashlib.md5(key_source.encode('utf8')).hexdigest()
 
 
-class ShelveCache(BudouCache):
+class PickleCache(BudouCache):
 
   def __init__(self, filepath = None):
     self.filepath = filepath if filepath else self.DEFAULT_FILE_PATH
 
   def get(self, source, language):
-    cache_shelve = shelve.open(self.filepath)
-    cache_key = self._get_cache_key(source, language)
-    result_value = cache_shelve.get(cache_key, None)
-    cache_shelve.close()
-    return result_value
+    self._create_file_if_none_exists()
+    with open(self.filepath, 'rb') as file_object:
+      cache_pickle = pickle.load(file_object)
+      cache_key = self._get_cache_key(source, language)
+      result_value = cache_pickle[cache_key]
+      return result_value
 
   def set(self, source, language, value):
-    cache_shelve = shelve.open(self.filepath)
-    cache_key = self._get_cache_key(source, language)
-    cache_shelve[cache_key] = value
-    cache_shelve.close()
+    self._create_file_if_none_exists()
+    with open(self.filepath, 'r+b') as file_object:
+      cache_pickle = pickle.load(file_object)
+      cache_key = self._get_cache_key(source, language)
+      cache_pickle[cache_key] = value
+      file_object.seek(0)
+      pickle.dump(cache_pickle, file_object)
+
+  def _create_file_if_none_exists(self):
+    if os.path.exists(self.filepath):
+      return
+    with open(self.filepath, 'wb') as file_object:
+      pickle.dump({}, file_object)
 
 
 class AppEngineCache(BudouCache):

test/cachefactory_test.py

@@ -30,21 +30,16 @@ def tearDown(self):
 
   def test_load(self):
     cache_type = repr(self.cache)
-    self.assertEqual('<ShelveCache>', cache_type,
-        'ShelveCache should be loaded under standard environment.')
+    self.assertEqual('<PickleCache>', cache_type,
+        'PickleCache should be loaded under standard environment.')
 
   def test_set_and_get(self):
     source = 'apple'
     language = 'a'
     target = 'banana'
     self.cache.set(source, language, target)
-
-
-    # Removing this test because shelve module can put a different extra suffix
-    # by low-level library.
-    # TODO (tushuhei) Update the cache mechanism to use pickle instead.
-    #self.assertTrue(os.path.isfile(self.cache.DEFAULT_FILE_PATH),
-    #    'Cache file should be generated.')
+    self.assertTrue(os.path.isfile(self.cache.DEFAULT_FILE_PATH),
+        'Cache file should be generated.')
     self.assertEqual(self.cache.get(source, language), target,
         'The target should be cached.')