Add cache decorator.

Author: Matthieu Ancellin

labelled_functions/decorators.py

@@ -4,6 +4,7 @@
 
 from labelled_functions.abstract import Unknown
 from labelled_functions.labels import label, LabelledFunction
+from copy import copy
 
 
 # API
@@ -49,22 +50,15 @@ def timed_func(*args, **kwargs):
     return timed_func
 
 
-# def cache(func, memory=None):
-#     """Return a copy of func but with results cached with joblib.
+def decorate(func, decorator):
+    new_func = copy(func)
+    new_func.function = decorator(func.function)
+    return new_func
 
-#     TODO: When used on a pipeline, do not merge the pipeline into a single function.
-#     """
-#     from labelled_functions.labels import LabelledFunction
 
-#     if memory is None:
-#         from joblib import Memory
-#         memory = Memory("/tmp", verbose=0)
-
-#     return LabelledFunction(
-#         memory.cache(func),
-#         name=func.name,
-#         input_names=func.input_names,
-#         output_names=func.output_names,
-#         default_values=func.default_values,
-#     )
+def cache(func, memory=None):
+    if memory is None:
+        from joblib import Memory
+        memory = Memory("/tmp", verbose=0)
+    return decorate(func, memory.cache)
 

test/test_decorators.py

@@ -3,7 +3,8 @@
 
 import pytest
 
-from labelled_functions.decorators import keeping_inputs
+from labelled_functions import label, pipeline
+from labelled_functions.decorators import keeping_inputs, decorate, cache
 from example_functions import *
 
 
@@ -26,3 +27,16 @@ def test_keeping_inputs():
     with pytest.raises(TypeError):
         keeping_inputs(all_kinds_of_args)(0, 1, 2, 3)
 
+
+def test_cache():
+    # Cache one function
+    pipe = cache(label(random_radius)) | label(cylinder_volume)
+    a = pipe(length=1.0)
+    b = pipe(length=1.0)
+    assert a == b
+
+    # Cache whole function
+    pipe = cache(pipeline([random_radius, cylinder_volume]))
+    a = pipe(length=1.0)
+    b = pipe(length=1.0)
+    assert a == b