diff --git a/atools/_memoize_decorator.py b/atools/_memoize_decorator.py index 31d36f3..ac7e47b 100644 --- a/atools/_memoize_decorator.py +++ b/atools/_memoize_decorator.py @@ -11,6 +11,7 @@ from time import time from threading import Lock as SyncLock from typing import Any, Callable, Dict, Hashable, Mapping, Optional, Tuple, Type, Union +from weakref import WeakSet Decoratee = Union[Callable, Type] @@ -492,7 +493,7 @@ def bar() -> Any: ... """ _default_db_path = Path.home() / '.memoize' - _all_decorators = set() + _all_decorators = WeakSet() @staticmethod def __call__( diff --git a/test/test_memoize_decorator.py b/test/test_memoize_decorator.py index 2212b02..d072da2 100644 --- a/test/test_memoize_decorator.py +++ b/test/test_memoize_decorator.py @@ -903,3 +903,19 @@ def foo(): assert foo() == PosixPath.cwd() assert isinstance(foo(), PosixPath) + + +def test_memoized_function_is_deletable() -> None: + def get_foo() -> Callable[[], None]: + @memoize + def _foo() -> None: + ... + + return _foo + + foo = get_foo() + r = ref(foo) + assert r() is not None + del foo + # FIXME there's a race condition here. Garbage collector may not have cleaned up foo yet + assert r() is None