As a toy example, let's use the Fibonacci sequence:
def fib(n: int) -> int:
if n < 2:
return 1
return fib(n - 2) + fib(n - 1)
Of course, this will hang the computer if we try to:
print(fib(100))
So we decide to add memoization. To keep the logic of fib clear, we decide not to change fib and instead add memoization via a decorator:
from typing import Callable
from functools import wraps
def remember(f: Callable[[int], int]) -> Callable[[int], int]:
@wraps(f)
def wrapper(n: int) -> int:
if n not in wrapper.memory:
wrapper.memory[n] = f(n)
return wrapper.memory[n]
wrapper.memory = dict[int, int]()
return wrapper
@remember
def fib(n: int) -> int:
if n < 2:
return 1
return fib(n - 2) + fib(n - 1)
Now there is no problem if we:
print(fib(100))
573147844013817084101
However, mypy complains that "Callable[[int], int]" has no attribute "memory", which makes sense, and usually I would want this complaint if I tried to access a property that is not part of the declared type...
So, how should we use typing to indicate that wrapper, while a Callable, also has the property memory?
To describe something as "a callable with a memory attribute", you could define a protocol (Python 3.8+, or earlier versions with typing_extensions):
from typing import Protocol
class Wrapper(Protocol):
memory: dict[int, int]
def __call__(self, n: int) -> int: ...
In use, the type checker knows that a Wrapper is valid as a Callable[[int], int] and allows return wrapper as well as the assignment to wrapper.memory:
from functools import wraps
from typing import Callable, cast
def remember(f: Callable[[int], int]) -> Callable[[int], int]:
@wraps(f)
def _wrapper(n: int) -> int:
if n not in wrapper.memory:
wrapper.memory[n] = f(n)
return wrapper.memory[n]
wrapper = cast(Wrapper, _wrapper)
wrapper.memory = dict()
return wrapper
Unfortunately this requires wrapper = cast(Wrapper, _wrapper), which is not type safe - wrapper = cast(Wrapper, "foo") would also check just fine.
Don't use a function attribute to store the cache, and you won't have this problem. You are already defining a closure (wrapper keeps a reference to the original callable), so store the cache in the closure as well.
from typing import Callable
from functools import wraps
def remember(f: Callable[[int], int]) -> Callable[[int], int]:
cache: dict[int, int] = {}
@wraps(f)
def wrapper(n: int) -> int:
if n not in cache:
cache[n] = f(n)
return cache[n]
return wrapper
@remember
def fib(n: int) -> int:
if n < 2:
return 1
return fib(n - 2) + fib(n - 1)
Building off of jonrsharpe's answer (which works, suggests the following, and I've accepted), we can avoid the need for a non type safe cast as follows:
from typing import Callable
from functools import wraps
class Remember:
def __init__(self) -> None:
self.memory = dict[int, int]()
def __call__(self, f: Callable[[int], int]) -> Callable[[int], int]:
@wraps(f)
def wrapper(n: int) -> int:
if n not in self.memory:
self.memory[n] = f(n)
return self.memory[n]
return wrapper
@Remember()
def fib(n: int) -> int:
if n < 2:
return 1
return fib(n - 2) + fib(n - 1)
print(fib(100))