Provide JSONDisk with diskcache (#124)

zsimic · grantjenks · commit 9cfe5df3fa0a · 2019-10-13T15:42:48.000-07:00
diff --git a/.gitignore b/.gitignore
@@ -3,6 +3,7 @@
 
 # virutalenv directories
 /env*/
+/.venv*/
 
 # test files/directories
 /.cache/
diff --git a/diskcache/__init__.py b/diskcache/__init__.py
@@ -6,7 +6,7 @@
 
 """
 
-from .core import Cache, Disk, EmptyDirWarning, UnknownFileWarning, Timeout
+from .core import Cache, Disk, EmptyDirWarning, JSONDisk, UnknownFileWarning, Timeout
 from .core import DEFAULT_SETTINGS, ENOVAL, EVICTION_POLICY, UNKNOWN
 from .fanout import FanoutCache
 from .persistent import Deque, Index
@@ -25,6 +25,7 @@
     'EmptyDirWarning',
     'FanoutCache',
     'Index',
+    "JSONDisk",
     'Lock',
     'RLock',
     'Timeout',
diff --git a/diskcache/core.py b/diskcache/core.py
@@ -7,6 +7,7 @@
 import errno
 import functools as ft
 import io
+import json
 import os
 import os.path as op
 import pickletools
@@ -361,6 +362,34 @@ def remove(self, filename):
                 raise
 
 
+class JSONDisk(Disk):
+    """Cache key and value (de)serialized as JSON."""
+    def __init__(self, directory, compress_level=1, **kwargs):
+        self.compress_level = compress_level
+        super(JSONDisk, self).__init__(directory, **kwargs)
+
+    def put(self, key):
+        json_bytes = json.dumps(key).encode('utf-8')
+        data = zlib.compress(json_bytes, self.compress_level)
+        return super(JSONDisk, self).put(data)
+
+    def get(self, key, raw):
+        data = super(JSONDisk, self).get(key, raw)
+        return json.loads(zlib.decompress(data).decode('utf-8'))
+
+    def store(self, value, read, key=UNKNOWN):
+        if not read:
+            json_bytes = json.dumps(value).encode('utf-8')
+            value = zlib.compress(json_bytes, self.compress_level)
+        return super(JSONDisk, self).store(value, read, key=key)
+
+    def fetch(self, mode, filename, value, read):
+        data = super(JSONDisk, self).fetch(mode, filename, value, read)
+        if not read:
+            data = json.loads(zlib.decompress(data).decode('utf-8'))
+        return data
+
+
 class Timeout(Exception):
     "Database timeout expired."
 
diff --git a/docs/tutorial.rst b/docs/tutorial.rst
@@ -813,7 +813,8 @@ database while values are sometimes stored separately in files.
 To customize serialization, you may pass in a :class:`Disk <diskcache.Disk>`
 subclass to initialize the cache. All clients accessing the cache are expected
 to use the same serialization. The default implementation uses Pickle and the
-example below uses compressed JSON.
+example below uses compressed JSON,
+available for convenience as :class:`Disk <diskcache.JSONDisk>`.
 
 .. code-block:: python
 
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -7,7 +7,6 @@
 import functools as ft
 import hashlib
 import io
-import json
 import mock
 import os
 import os.path as op
@@ -22,7 +21,6 @@
 import time
 import unittest
 import warnings
-import zlib
 
 try:
     import cPickle as pickle
@@ -94,35 +92,8 @@ def test_disk_valueerror():
             pass
 
 
-class JSONDisk(diskcache.Disk):
-    def __init__(self, directory, compress_level=1, **kwargs):
-        self.compress_level = compress_level
-        super(JSONDisk, self).__init__(directory, **kwargs)
-
-    def put(self, key):
-        json_bytes = json.dumps(key).encode('utf-8')
-        data = zlib.compress(json_bytes, self.compress_level)
-        return super(JSONDisk, self).put(data)
-
-    def get(self, key, raw):
-        data = super(JSONDisk, self).get(key, raw)
-        return json.loads(zlib.decompress(data).decode('utf-8'))
-
-    def store(self, value, read, key=dc.UNKNOWN):
-        if not read:
-            json_bytes = json.dumps(value).encode('utf-8')
-            value = zlib.compress(json_bytes, self.compress_level)
-        return super(JSONDisk, self).store(value, read, key=key)
-
-    def fetch(self, mode, filename, value, read):
-        data = super(JSONDisk, self).fetch(mode, filename, value, read)
-        if not read:
-            data = json.loads(zlib.decompress(data).decode('utf-8'))
-        return data
-
-
 def test_custom_disk():
-    with dc.Cache(disk=JSONDisk, disk_compress_level=6) as cache:
+    with dc.Cache(disk=dc.JSONDisk, disk_compress_level=6) as cache:
         values = [None, True, 0, 1.23, {}, [None] * 10000]
 
         for value in values:
