使用mypy为Python实例方法编写装饰器时如何访问实例属性？

问题背景 #

在Python项目中使用mypy做静态类型检查时，为Database类的实例方法编写require_connection装饰器，需要确保调用方法前数据库已完成连接（即self._connection不为None），但原装饰器无法正确访问实例的self._connection属性，同时需要保证类型标注能通过mypy检查。

原Database类代码：

import sqlite3
from typing import ParamSpec, TypeVar, Callable, wraps

# 假设InputTable是已定义的类型
class InputTable:
    pass

class DataBase:
    _path: str
    _connection: sqlite3.Connection | None
    _cursor: sqlite3.Cursor | None

    def __init__(self, path: str) -> None:
        self._path = path
        self._connection = None
        self._cursor = None

    def connect(self) -> None:
        self._connection = sqlite3.connect(
            self._path, detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES
        )
        self._connection.execute("PRAGMA foreign_keys = ON;")
        self._cursor = self._connection.cursor()

    @require_connection
    def add_input_data(self, inputData: InputTable) -> None:
        # 业务逻辑
        pass

    @require_connection
    def close(self) -> None:
        if self._connection:
            self._connection.close()
            self._connection = None
            self._cursor = None

原装饰器代码（存在问题）：

P = ParamSpec("P")
T = TypeVar("T")

def require_connection(func: Callable[P, T]) -> Callable[P, T]:
    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        # 无法访问self，mypy会报错
        if self._connection is None:
            raise Exception(
                f"Cannot use the method: {func.__name__} before connecting to the Database"
            )
        return func(*args, **kwargs)

    return wrapper

解决步骤 #

1. 从参数中提取实例self #

实例方法调用时，第一个参数必然是类的实例对象（即self），因此可以从args的第一个元素中提取，并通过类型断言让mypy识别其类型。

2. 修正装饰器的类型标注 #

使用绑定到DataBase类的TypeVar，结合ParamSpec确保装饰器仅适用于DataBase的实例方法，同时保留原方法的参数和返回值类型。

修正后的装饰器代码 #

from typing import ParamSpec, TypeVar, Callable, wraps
import sqlite3

# 定义绑定到DataBase的TypeVar
DBInstance = TypeVar("DBInstance", bound="DataBase")
P = ParamSpec("P")
T = TypeVar("T")

def require_connection(func: Callable[P, T]) -> Callable[P, T]:
    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        # 提取第一个参数作为self实例
        self = args[0]
        # 类型断言，让mypy识别为DataBase实例
        if not isinstance(self, DataBase):
            raise TypeError("Decorator can only be used on DataBase instance methods")
        
        if self._connection is None:
            raise Exception(
                f"Cannot use method {func.__name__} before connecting to the database"
            )
        return func(*args, **kwargs)

    return wrapper

完整可运行代码 #

import sqlite3
from typing import ParamSpec, TypeVar, Callable, wraps

class InputTable:
    pass

DBInstance = TypeVar("DBInstance", bound="DataBase")
P = ParamSpec("P")
T = TypeVar("T")

def require_connection(func: Callable[P, T]) -> Callable[P, T]:
    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        self = args[0]
        if not isinstance(self, DataBase):
            raise TypeError("Decorator can only be used on DataBase instance methods")
        
        if self._connection is None:
            raise Exception(
                f"Cannot use method {func.__name__} before connecting to the database"
            )
        return func(*args, **kwargs)

    return wrapper

class DataBase:
    _path: str
    _connection: sqlite3.Connection | None
    _cursor: sqlite3.Cursor | None

    def __init__(self, path: str) -> None:
        self._path = path
        self._connection = None
        self._cursor = None

    def connect(self) -> None:
        self._connection = sqlite3.connect(
            self._path, detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES
        )
        self._connection.execute("PRAGMA foreign_keys = ON;")
        self._cursor = self._connection.cursor()

    @require_connection
    def add_input_data(self, inputData: InputTable) -> None:
        # 示例业务逻辑
        if self._cursor:
            self._cursor.execute("INSERT INTO input_data VALUES (?)", (inputData,))
            self._connection.commit()

    @require_connection
    def close(self) -> None:
        if self._connection:
            self._connection.close()
            self._connection = None
            self._cursor = None

关键说明 #

• 通过args[0]获取实例对象，解决装饰器中无法直接访问self的问题；
• 加入isinstance检查和类型断言，让mypy能够正确推断实例类型，避免类型错误；
• 使用bound参数的TypeVar限制装饰器仅用于DataBase类的实例方法，提升类型安全性。

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