Add pickler argument to memoize.

.travis.yml

@@ -7,7 +7,7 @@ dist: xenial
 env:
   PYTHONPATH=.
 before_install:
-  - pip install pytest pytest-asyncio pytest-cov
+  - pip install dill pytest pytest-asyncio pytest-cov
   - pip install coveralls
 install:
   - python setup.py -q install

README.md

@@ -11,6 +11,7 @@ Decorates a function call and caches return value for given inputs.
 - If `db_path` is provided, memos will persist on disk and reloaded during initialization.
 - If `duration` is provided, memos will only be valid for given `duration`.
 - If `keygen` is provided, memo hash keys will be created with given `keygen`.
+- If `pickler` is provided, persistent memos will (de)serialize using given `pickler`.
 - If `size` is provided, LRU memo will be evicted if current count exceeds given `size`.
 
 ### Examples
@@ -263,6 +264,16 @@ inputs.
     foo.bar()  # Function not called. Cached result returned.
     ```
 
+- Custom pickler may be specified for persistent memo (de)serialization.
+
+    ```python3
+    import dill
+
+    @memoize(db_path='~/.memoize`, pickler=dill)
+    def foo() -> Callable[[], None]:
+        return lambda: None
+    ```
+
 ## rate
 Function decorator that rate limits the number of calls to function.
 

atools/_memoize_decorator.py

@@ -1,3 +1,4 @@
+from abc import ABC
 from asyncio import Lock as AsyncLock
 from collections import ChainMap, OrderedDict
 from dataclasses import dataclass, field
@@ -11,14 +12,25 @@
 from textwrap import dedent
 from time import time
 from threading import Lock as SyncLock
-from typing import Any, Callable, Iterable, Hashable, Mapping, Optional, Tuple, Type, Union
+from typing import Any, Callable, Hashable, Mapping, Optional, Tuple, Type, Union
 from weakref import finalize, WeakSet
 
 
 Decoratee = Union[Callable, Type]
 Keygen = Callable[..., Any]
 
 
+class Pickler(ABC):
+
+    @staticmethod
+    def dumps(_str: str) -> str:
+        ...  # pragma: no cover
+
+    @staticmethod
+    def loads(_bytes: bytes) -> Any:
+        ...  # pragma: no cover
+
+
 class _MemoZeroValue:
     pass
 
@@ -56,6 +68,7 @@ class _MemoizeBase:
     duration: Optional[timedelta]
     fn: Callable
     keygen: Optional[Keygen]
+    pickler: Pickler = field(hash=False)
     size: Optional[int]
 
     expire_order: OrderedDict = field(init=False, default_factory=OrderedDict, hash=False)
@@ -91,7 +104,7 @@ def __post_init__(self) -> None:
             ).fetchall():
                 memo = self.make_memo(t0=t0)
                 memo.memo_return_state.called = True
-                memo.memo_return_state.value = pickle.loads(v)
+                memo.memo_return_state.value = self.pickler.loads(v)
                 self.memos[k] = memo
             if self.duration:
                 for k, t0 in self.db.execute(
@@ -168,7 +181,7 @@ def finalize_memo(self, memo: _Memo, key: Union[int, str]) -> Any:
         if memo.memo_return_state.raised:
             raise memo.memo_return_state.value
         elif (self.db is not None) and (self.memos[key] is memo):
-            value = pickle.dumps(memo.memo_return_state.value)
+            value = self.pickler.dumps(memo.memo_return_state.value)
             self.db.execute(
                 dedent(f'''
                     INSERT OR REPLACE INTO `{self.table_name}`
@@ -413,6 +426,7 @@ class _Memoize:
     - If `db_path` is provided, memos will persist on disk and reloaded during initialization.
     - If `duration` is provided, memos will only be valid for given `duration`.
     - If `keygen` is provided, memo hash keys will be created with given `keygen`.
+    - If `pickler` is provided, persistent memos will (de)serialize using given `pickler`.
     - If `size` is provided, LRU memo will be evicted if current count exceeds given `size`.
 
     ### Examples
@@ -664,6 +678,16 @@ def bar(self) -> Any: ...
         # foo instance is kept around somewhere and used later.
         foo.bar()  # Function not called. Cached result returned.
         ```
+
+    - Custom pickler may be specified for unpickleable return types.
+
+        ```python3
+        import dill
+
+        @memoize(db_path='~/.memoize`, pickler=dill)
+        def foo() -> Callable[[], None]:
+            return lambda: None
+        ```
     """
 
     _all_decorators = WeakSet()
@@ -675,10 +699,11 @@ def __call__(
             db_path: Optional[Path] = None,
             duration: Optional[Union[int, float, timedelta]] = None,
             keygen: Optional[Keygen] = None,
+            pickler: Optional[Pickler] = None,
             size: Optional[int] = None,
     ) -> Union[Decoratee]:
         if _decoratee is None:
-            return partial(memoize, db_path=db_path, duration=duration, keygen=keygen, size=size)
+            return partial(memoize, db_path=db_path, duration=duration, keygen=keygen, pickler=pickler, size=size)
 
         if inspect.isclass(_decoratee):
             assert db_path is None, 'Class memoization not allowed with db.'
@@ -697,6 +722,7 @@ class Wrapped(_decoratee, metaclass=WrappedMeta):
         db = connect(f'{db_path}') if db_path is not None else None
         duration = timedelta(seconds=duration) if isinstance(duration, (int, float)) else duration
         assert (duration is None) or (duration.total_seconds() > 0)
+        pickler = pickle if pickler is None else pickler
         assert (size is None) or (size > 0)
         fn = _decoratee
         default_kwargs: Mapping[str, Any] = {
@@ -715,6 +741,7 @@ class Wrapped(_decoratee, metaclass=WrappedMeta):
             duration=duration,
             fn=fn,
             keygen=keygen,
+            pickler=pickler,
             size=size,
         ).get_decorator()
 

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='atools',
-    version='0.13.1',
+    version='0.14.0',
     packages=find_packages(),
     python_requires='>=3.6',
     url='https://github.com/cevans87/atools',
@@ -11,6 +11,7 @@
     author_email='c.d.evans87@gmail.com',
     description='Python 3.6+ async/sync memoize and rate decorators',
     tests_require=[
+        'dill',
         'pytest',
         'pytest-asyncio',
         'pytest-cov',

test/test_memoize_decorator.py

@@ -31,6 +31,13 @@ def db_path() -> Path:
         yield Path(f.name)
 
 
+@pytest.fixture
+def dill() -> test_module.Pickler:
+    import dill
+    yield dill
+    del dill
+
+
 @pytest.fixture
 def sync_lock() -> MagicMock:
     with patch.object(test_module, 'SyncLock', side_effect=None) as sync_lock:
@@ -1158,3 +1165,18 @@ async def foo() -> None:
     assert body.call_count == 1
     assert len(foo.memoize) == 1
     assert await foo() == 1
+
+
+def test_function_return_type_with_db_and_dill_does_not_raise(db_path: Path, dill: test_module.Pickler) -> None:
+    foo_body = MagicMock()
+
+    @memoize(db_path=db_path, pickler=dill)
+    def foo() -> Callable[[], None]:
+        foo_body()
+
+        return lambda: None
+
+    foo()()
+    foo()()
+
+    assert foo_body.call_count == 1