如何让装饰器收窄函数体内类型，使mypy识别非空约束？

要让mypy识别装饰器已经确保self.data不为None，可以结合TypeGuard和Protocol给装饰器添加精确的类型标注，以下是具体实现方案：

步骤1：定义类型协议 #

通过Protocol明确区分data为None和非None的实例类型，让mypy能识别两种状态：

from typing import Protocol, TypeGuard, Callable, TypeVar, Exclude
from functools import wraps

class HasNoData(Protocol):
    data: None
    def log_something(self) -> None: ...

class HasData(Protocol):
    data: list[int]

步骤2：实现带TypeGuard的检查函数 #

编写提前退出的判断函数，用TypeGuard告诉mypy：当函数返回True时，实例属于HasNoData类型（即data为None）：

def data_is_none(obj: HasData | HasNoData) -> TypeGuard[HasNoData]:
    if obj.data is None:
        obj.log_something()
        return True
    return False

步骤3：实现泛型装饰器 #

给装饰器添加泛型类型标注，让mypy能推断出：经过装饰器的提前退出后，进入方法体的实例已经排除了HasNoData类型，必然是HasData：

T = TypeVar('T')
U = TypeVar('U')

def early_exit_if(cond: Callable[[T], TypeGuard[U]]) -> Callable[[Callable[[Exclude[T, U]], None]], Callable[[T], None]]:
    def decorator(func: Callable[[Exclude[T, U]], None]) -> Callable[[T], None]:
        @wraps(func)
        def wrapper(self: T) -> None:
            if cond(self):
                return
            func(self)
        return wrapper
    return decorator

步骤4：使用装饰器 #

现在在类方法中使用装饰器，mypy会自动识别方法体内self.data不为None：

class MyClass:
    def __init__(self):
        self.data: list[int] | None = None

    def log_something(self) -> None:
        print("Data is None, exiting early")

    @early_exit_if(data_is_none)
    def do_something_with_data(self) -> None:
        self.data.append(1)  # mypy不再报错：能识别self.data是list[int]

关键说明 #

• Exclude[T, U]用于从类型T中排除U，这里就是从HasData | HasNoData中排除HasNoData，剩下HasData类型。
• 确保mypy启用了strict或strict_protocols配置，这样协议的类型检查会更准确。
• 这种方式不需要在方法体内添加assert self.data is not None，完全通过装饰器的类型标注实现类型收窄。

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