Repository: cevans87/atools

Issue:
Use function return type-hint in `eval` when memoize deserializes from db.
For example
```
@memoize(db=True)
def foo() -> pathlib.Path:
    ...
```
This cannot currently be stored and deserialized from the memoize db since we currently only allow deserialization of `__builtins__`. Allowing `__builtins__` _plus_ the return type-hint will greatly expand the usability.

Discussion:

Patch:
diff --git a/atools/_memoize_decorator.py b/atools/_memoize_decorator.py
index 783caf2..3b0c7ab 100644
--- a/atools/_memoize_decorator.py
+++ b/atools/_memoize_decorator.py
@@ -10,7 +10,7 @@ from sqlite3 import connect, Connection
 from textwrap import dedent
 from time import time
 from threading import Lock as SyncLock
-from typing import Any, Callable, Hashable, Mapping, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, Hashable, Mapping, Optional, Tuple, Type, Union
 
 
 Decoratee = Union[Callable, Type]
@@ -57,11 +57,22 @@ class _MemoizeBase:
     keygen: Optional[Keygen]
     size: Optional[int]
 
+    db_eval_types: Optional[Dict[str, Any]] = field(init=False, default=None, hash=False)
     expire_order: OrderedDict = field(init=False, default_factory=OrderedDict, hash=False)
     memos: OrderedDict = field(init=False, default_factory=OrderedDict, hash=False)
 
     def __post_init__(self) -> None:
         if self.db is not None:
+            return_annotation = inspect.signature(self.fn).return_annotation
+            if not hasattr(return_annotation, '__name__'):
+                object.__setattr__(self, 'db_eval_types', dict(self.fn.__globals__))
+            else:
+                object.__setattr__(
+                    self,
+                    'db_eval_types',
+                    {**self.fn.__globals__, return_annotation.__name__: return_annotation}
+                )
+
             self.db.execute(dedent(f'''
                 CREATE TABLE IF NOT EXISTS `{self.table_name}` (
                   k TEXT PRIMARY KEY,
@@ -88,7 +99,7 @@ class _MemoizeBase:
             ).fetchall():
                 memo = self.make_memo(fn=self.fn, t0=t0)
                 memo.memo_return_state.called = True
-                (memo.memo_return_state.value,) = eval(v, __builtins__)
+                (memo.memo_return_state.value,) = eval(v, self.db_eval_types)
                 self.memos[k] = memo
             if self.duration:
                 for k, t0 in self.db.execute(
@@ -159,7 +170,7 @@ class _MemoizeBase:
         else:
             if (self.db is not None) and (self.memos[key] is memo):
                 value = str((memo.memo_return_state.value,))
-                assert (memo.memo_return_state.value,) == eval(value, __builtins__)
+                assert (memo.memo_return_state.value,) == eval(value, self.db_eval_types)
                 self.db.execute(
                     dedent(f'''
                         INSERT OR REPLACE INTO `{self.table_name}`
diff --git a/setup.py b/setup.py
index 45f21de..723457b 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import find_packages, setup
 
 setup(
     name='atools',
-    version='0.10.2',
+    version='0.10.3',
     packages=find_packages(),
     python_requires='>=3.6',
     url='https://github.com/cevans87/atools',


Test patch:
diff --git a/test/test_memoize_decorator.py b/test/test_memoize_decorator.py
index b535cf6..2212b02 100644
--- a/test/test_memoize_decorator.py
+++ b/test/test_memoize_decorator.py
@@ -4,7 +4,7 @@ from asyncio import (
 from atools import memoize
 import atools._memoize_decorator as test_module
 from datetime import timedelta
-from pathlib import Path
+from pathlib import Path, PosixPath
 import pytest
 from sqlite3 import connect, Connection
 from tempfile import NamedTemporaryFile
@@ -572,6 +572,7 @@ async def test_memoize_does_not_stop_object_cleanup() -> None:
     r = ref(f)
     assert r() is not None
     del f
+    # FIXME there's a race condition here. Garbage collector may not have cleaned up f yet
     assert r() is None
 
 
@@ -860,7 +861,7 @@ def test_set_default_db_path_uses_given_path() -> None:
         assert get_table_len(f.name) == 0
 
         @memoize(db=True)
-        def foo():
+        def foo() -> None:
             ...
 
         assert get_table_len(f.name) == 1
@@ -874,9 +875,31 @@ async def test_async_keygen_can_return_non_tuple() -> None:
         return 1
 
     @memoize(keygen=lambda: keygen())
-    async def foo():
+    async def foo() -> None:
         body()
 
     await foo()
     await foo()
     assert body.call_count == 1
+
+
+def test_db_can_return_type_of_return_hint(db: Union[bool, Connection, Path, str]) -> None:
+    class FooPath(PosixPath):
+        pass
+
+    @memoize(db=db)
+    def foo() -> FooPath:
+        return FooPath.cwd()
+
+    assert foo() == FooPath.cwd()
+    assert isinstance(foo(), FooPath)
+
+
+def test_db_can_return_type_of_callers_globals(db: Union[bool, Connection, Path, str]) -> None:
+
+    @memoize(db=db)
+    def foo():
+        return PosixPath.cwd()
+
+    assert foo() == PosixPath.cwd()
+    assert isinstance(foo(), PosixPath)