通过装饰器实现Memoization时，如何保留原函数的Type hints？

要让Memoize装饰器既保持通用性，又能让类型检查器自动适配原函数的返回类型，你可以用Python的**泛型（Generic）**来改造它，具体实现如下：

from typing import TypeVar, Generic, Callable, Any

T = TypeVar('T')  # 定义泛型占位符，用来指代任意返回类型

class A:
    def __init__(self, f):
        self.foo = 1337

class Memoize(Generic[T]):
    def __init__(self, f: Callable[..., T]):
        self.f = f
        self.memo = {}
    
    def __call__(self, *args: Any) -> T:
        if args not in self.memo:
            self.memo[args] = self.f(*args)
        return self.memo[args]

@Memoize
def get_file(file_path: str) -> A:
    return A(file_path)

a = get_file("bar.txt").foo  # VS Code 现在可以正确解析a的类型

关键说明 #

• TypeVar('T')：创建一个类型变量，用来关联原函数的返回类型，让装饰器可以适配任意返回类型的函数。
• Memoize(Generic[T])：让装饰器类继承泛型类，使其能和原函数的返回类型T绑定。
• Callable[..., T]：指定传入的函数f是一个可调用对象，无论参数是什么，返回类型都是T。
• __call__方法返回类型标注为T：确保装饰后的函数返回类型和原函数完全一致，类型检查器就能自动识别。

这样改造后，Memoize装饰器依然可以用来缓存任意函数的结果，同时VS Code等工具能正确解析被装饰函数的返回类型，包括get_file返回的A类型，以及a的类型。

内容的提问来源于stack exchange，提问作者Sandro