How should we type a callable with additional properties?

Question

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?

Answer 1

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

Playground

Unfortunately this requires wrapper = cast(Wrapper, _wrapper), which is not type safe - wrapper = cast(Wrapper, "foo") would also check just fine.

Answer 2

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)