Add diskcache

Author: danijar

elements/__init__.py

@@ -1,9 +1,10 @@
-__version__ = '3.19.1'
+__version__ = '3.20.0'
 
 from .agg import Agg
 from .checkpoint import Checkpoint, Saveable
 from .config import Config
 from .counter import Counter
+from .diskcache import diskcache
 from .flags import Flags
 from .fps import FPS
 from .logger import Logger

elements/diskcache.py

@@ -0,0 +1,64 @@
+import hashlib
+import json
+import pickle
+
+from . import path as pathlib
+
+
+def diskcache(*args, **kwargs):
+  if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
+    return DiskCache(args[0])
+  else:
+    return lambda fn: DiskCache(fn, *args, **kwargs)
+
+
+diskcache.root = '/tmp/elements/diskcache'
+diskcache.refresh = False
+diskcache.verbose = False
+
+
+class DiskCache:
+
+  def __init__(self, fn, name=None, refresh=None, verbose=None, root=None):
+    if name is None:
+      name = f'{pathlib.Path(__file__).stem}-{fn.__name__}'
+    self.fn = fn
+    self.folder = pathlib.Path(diskcache.root if root is None else root) / name
+    self.refresh = diskcache.refresh if refresh is None else refresh
+    self.verbose = diskcache.verbose if verbose is None else verbose
+
+  def __call__(
+      self,
+      *args,
+      _key=None,
+      _refresh=None,
+      **kwargs,
+  ):
+    try:
+      inputs = [args, kwargs] if _key is None else _key
+      inputs = json.dumps(inputs, sort_keys=True)
+    except ValueError as e:
+      raise ValueError(
+          'Diskcache requires function arguments to be JSON serializable. ' +
+          'Alternatively, pass _key=... into the function with a cache key.'
+      ) from e
+    key = hashlib.sha256(inputs.encode('utf8')).hexdigest()
+    filename = self.folder / f'{key}.pkl'
+    refresh = self.refresh if _refresh is None else _refresh
+    if filename.exists() and not refresh:
+      self.verbose and print(f'Loading diskcache: {filename}')
+      data = pickle.loads(filename.read_bytes())
+      assert data['inputs'] == inputs, ('Hash collision', data, inputs)
+      return data['output']
+    else:
+      self.verbose and print(f'Filling diskcache: {filename}')
+      output = self.fn(*args, **kwargs)
+      data = {'inputs': inputs, 'output': output}
+      filename.parent.mkdir()
+      filename.write_bytes(pickle.dumps(data))
+    return output
+
+  def clear(self):
+    self.verbose and print(f'Clearing diskcache: {self.folder}')
+    for file in self.folder.glob('*.pkl'):
+      file.remove()

tests/test_diskcache.py

@@ -0,0 +1,60 @@
+import elements
+import elements.diskcache
+
+
+class TestDiskcache:
+
+  def test_basic(self, tmpdir):
+    elements.diskcache.root = tmpdir
+    elements.diskcache.verbose = True
+
+    counter = elements.Counter()
+
+    @elements.diskcache
+    def fn(foo, bar):
+      counter.increment()
+      return foo + bar
+
+    fn.clear()
+
+    assert counter == 0
+    assert fn(1, 2) == 3
+    assert counter == 1
+    assert fn('a', 'b') == 'ab'
+    assert counter == 2
+
+    assert fn(1, 2) == 3
+    assert fn('a', 'b') == 'ab'
+    assert counter == 2
+
+    fn.clear()
+    assert fn(1, 2) == 3
+    assert fn('a', 'b') == 'ab'
+    assert counter == 4
+
+    assert fn(1, 2, _refresh=True) == 3
+    assert counter == 5
+
+    assert fn(1, 2) == 3
+    assert fn('a', 'b') == 'ab'
+    assert counter == 5
+
+
+  def test_names(self, tmpdir):
+
+    def make():
+      @elements.diskcache('fn1')
+      def fn(foo, bar):
+        return foo + bar
+      return fn
+    fn1 = make()
+
+    def make():
+      @elements.diskcache('fn2')
+      def fn(foo, bar):
+        return foo - bar
+      return fn
+    fn2 = make()
+
+    assert fn1(2, 1) == 3
+    assert fn2(2, 1) == 1