在 Python 编程中，装饰器是一种强大的工具，用于动态地修改函数的行为。随着类型注解在 Python 中的普及，如何在装饰器中精确地使用类型注解成为了一个热门话题。本文将探讨如何通过类型注解优化装饰器的使用，确保代码的类型安全性。

装饰器的类型注解基础

首先，我们来看一个简单的装饰器示例，它用于注册操作符：

from typing import Protocol, TypeVar, Generic, Callable

TIn = TypeVar('TIn', contravariant=True)
TOut = TypeVar('TOut', covariant=True)
IntFunction = Callable[[int, int], int]

class Decorator(Protocol, Generic[TIn, TOut]):
    def __call__(self, value: TIn) -> TOut: ...

def register_operator(op: str) -> Decorator[IntFunction, IntFunction]:
    def inner(value: IntFunction) -> IntFunction:
        # 注册函数或其他操作
        return value
    return inner

@register_operator("+")
def add(a: int, b: int) -> int:
    return a + b

在这个例子中，register_operator 装饰器接受一个字符串参数 op，并返回一个 Decorator 类型的函数，该函数接收并返回 IntFunction 类型。

优化装饰器的类型注解

如果我们希望简化 register_operator 的类型注解，使得输入和输出的类型可以默认相同，我们可以采用以下方法：

使用两个协议类

由于 Python 不支持泛型的默认值，我们可以创建两个协议类，其中一个表示对称的输入输出类型：

from typing import Protocol, TypeVar, Generic, overload

TIn = TypeVar('TIn', contravariant=True)
TOut = TypeVar('TOut', covariant=True)
TSym = TypeVar('TSym')

class Decorator(Protocol, Generic[TIn, TOut]):
    def __call__(self, value: TIn) -> TOut: ...

class SymmetricDecorator(Decorator[TSym, TSym], Generic[TSym], Protocol):
    pass

def register_operator(op: str) -> SymmetricDecorator[IntFunction]:
    def inner(value: IntFunction) -> IntFunction:
        # 注册函数或其他操作
        return value
    return inner

@register_operator("+")
def add(a: int, b: int) -> int:
    return a + b

在这里，我们定义了 SymmetricDecorator，它继承自 Decorator 并强制输入和输出类型相同。这样，我们在使用时可以简化类型注解：

def register_operator(op: str) -> SymmetricDecorator[IntFunction]:
    ...

实例分析

考虑一个实际的例子，我们可能有一个计算器类，其中每个操作符都需要通过装饰器注册：

class Calculator:
    def __init__(self):
        self.operators = {}

    @register_operator("+")
    def add(self, a: int, b: int) -> int:
        return a + b

    @register_operator("-")
    def subtract(self, a: int, b: int) -> int:
        return a - b

calc = Calculator()
calc.operators["+"](2, 3)  # 结果应该是 5

在这个例子中，我们使用了优化后的 register_operator 装饰器，它简化了类型注解的使用，同时保留了类型安全性。每个操作符函数都被注册到 Calculator 实例的 operators 字典中，以便后续调用。

结论

通过使用两个协议类，我们成功地简化了装饰器的类型注解，保持了类型安全性，并使代码更加清晰易读。随着 Python 类型系统的不断完善，未来或许会有更直接的支持默认值的泛型类型，但目前这种方法已经足以应对大部分的实际需求。

希望这篇博客对你理解和优化 Python 装饰器的类型注解有所帮助，欢迎在评论区分享你的见解或问题！