如何使MyPy识别校验类连接属性的Python装饰器类型断言

问题背景 #

原代码通过在每个方法内手动校验_connection是否为None，能通过MyPy类型检查；但改用requires_connection装饰器集中处理校验后，MyPy无法识别装饰器已确保_connection非空，抛出以下错误：

error: Item "None" of "Connection | None" has no attribute "read"  [union-attr]
error: Item "None" of "Connection | None" has no attribute "write"  [union-attr]

解决方案 #

方案一：在装饰器内断言并重新赋值连接变量 #

修改装饰器代码，将cast后的连接变量重新赋值给self._connection，让MyPy识别此时变量已非空：

from functools import wraps
from typing import cast

class Connection:
    def read(self) -> str:
        return "data"
    
    def write(self, data: str) -> None:
        print(data)

def requires_connection(func):
    @wraps(func)
    def decorator(self, *args, **kwargs):
        if self._connection is None:
            raise RuntimeError("Connection not created")
        # 断言类型并重新赋值，告知MyPy此时_connection非空
        self._connection = cast(Connection, self._connection)
        return func(self, *args, **kwargs)
    return decorator

class Foo:
    def __init__(self):
        self._connection: None | Connection = None
        
    def connect(self):
        self._connection = Connection()
        
    @requires_connection
    def get_data(self) -> str:
        return self._connection.read()
        
    @requires_connection
    def send_data(self, data: str):
        self._connection.write(data)

方案二：使用类型守卫（TypeGuard）定义校验逻辑 #

在Foo类中定义一个类型守卫方法，装饰器调用该方法后，MyPy会自动推断_connection的非空状态：

from functools import wraps
from typing import TypeGuard
# Python 3.11+可直接用typing.Self，低版本需从typing_extensions导入
from typing_extensions import Self

class Connection:
    def read(self) -> str:
        return "data"
    
    def write(self, data: str) -> None:
        print(data)

class Foo:
    def __init__(self):
        self._connection: None | Connection = None
        
    def connect(self):
        self._connection = Connection()
    
    # 类型守卫方法：校验通过则断言实例的_connection非空
    def _has_valid_connection(self) -> TypeGuard[Self]:
        if self._connection is None:
            raise RuntimeError("Connection not created")
        return True

def requires_connection(func):
    @wraps(func)
    def decorator(self: Foo, *args, **kwargs):
        self._has_valid_connection()
        return func(self, *args, **kwargs)
    return decorator

class Foo:
    # ... 其余代码不变 ...
    @requires_connection
    def get_data(self) -> str:
        return self._connection.read()
        
    @requires_connection
    def send_data(self, data: str):
        self._connection.write(data)

说明 #

• 方案一实现简单，仅需修改装饰器代码，适合快速解决问题；
• 方案二更符合类型系统规范，类型守卫的逻辑可复用，适合复杂类结构或多场景校验需求。

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