使用描述符和Pyright实现Python运算符重载的最佳实践

from typing import Callable as Fn, Any, overload
import operator

class Apply:
    """
    封装具体运算符逻辑的可调用对象。
    它接收一个操作函数和一个对象，并在被调用时应用该操作。
    """
    def __init__(self, op: Fn[[Any, Any], Any], obj: Any) -> None:
        self.op = op
        self.obj = obj

    # 模拟两个重载签名，实际应用中可以根据需要扩展
    @overload
    def __call__(self, x: int) -> str: ...
    @overload
    def __call__(self, x: str) -> int: ...

    def __call__(self, x: int | str) -> str | int:
        # 实际的运算符逻辑可以根据op和x的类型进行分发
        # 这里的实现仅为示例，实际应根据需求进行具体操作
        if isinstance(x, int):
            return str(self.op(self.obj, x))
        elif isinstance(x, str):
            return int(self.op(self.obj, int(x)))
        raise TypeError("Unsupported operand type")

class Op:
    """
    运算符的数据模型描述符。
    当通过实例访问时，它返回一个Apply对象，
    该对象封装了具体的运算符和所属实例。
    """
    def __init__(self, op: Fn[[Any, Any], Any]) -> None:
        self.op = op

    def __get__(self, obj: Any, _: Any) -> Apply:
        # 当通过实例（obj）访问时，返回一个绑定了op和obj的Apply实例
        return Apply(self.op, obj)

class Foo:
    # 将__add__和__mul__魔术方法设置为Op描述符的实例
    __add__ = Op(operator.add)
    __mul__ = Op(operator.mul)

# 实例化Foo类
foo = Foo()

# 此时，直接调用描述符返回的Apply对象是有效的，并且类型检查器可以正确推断
a: str = foo.__add__(2)    # Pyright: str (正确)
b: int = foo.__mul__("2")  # Pyright: int (正确)

# 然而，当尝试使用Python的语法糖直接应用运算符时，Pyright会报告类型错误
_ = foo + 1    # Pyright: 类型错误 (Expected type 'int', got 'Foo')
_ = foo * "2"  # Pyright: 类型错误 (Expected type 'str', got 'Foo')

class Op:
    """
    运算符的数据模型描述符。
    当通过实例访问时，它返回一个Apply对象，
    该对象封装了具体的运算符和所属实例。
    """
    def __init__(self, op: Fn[[Any, Any], Any]) -> None:
        self.op = op

    def __get__(self, obj: Any, _: Any) -> Apply:
        return Apply(self.op, obj)

    # 关键的辅助类型注解：
    # 明确告诉Pyright，Op实例（作为Foo类的__add__等属性）
    # 在被调用时，其行为与Apply对象相同。
    __call__: Apply

from typing import Callable as Fn, Any, overload
import operator

class Apply:
    """Apply an operator to an object."""
    def __init__(self, op: Fn[[Any, Any], Any], obj: Any) -> None:
        self.op = op
        self.obj = obj

    @overload
    def __call__(self, x: int) -> str: ...
    @overload
    def __call__(self, x: str) -> int: ...

    def __call__(self, x: int | str) -> str | int:
        if isinstance(x, int):
            return str(self.op(self.obj, x))
        elif isinstance(x, str):
            return int(self.op(self.obj, int(x)))
        raise TypeError("Unsupported operand type")

class Op:
    """Data model object for an operator."""
    def __init__(self, op: Fn[[Any, Any], Any]) -> None:
        self.op = op

    def __get__(self, obj: Any, _: Any) -> Apply:
        return Apply(self.op, obj)

    __call__: Apply  # 辅助注解

class Foo:
    __add__ = Op(operator.add)
    __mul__ = Op(operator.mul)

foo = Foo()

# 使用reveal_type来查看Pyright的类型推断结果
# (在实际运行中，reveal_type需要pyright命令行工具支持)
# reveal_type(foo.__add__(2))    # 预期: str
# reveal_type(foo.__mul__("2"))  # 预期: int
# reveal_type(foo + 1)           # 预期: str
# reveal_type(foo * "2")         # 预期: int

# 实际测试结果（Pyright playground验证）：
a: str = foo.__add__(2)    # Pyright: str (正确)
b: int = foo.__mul__("2")  # Pyright: int (正确)
c: str = foo + 1           # Pyright: str (正确，不再报错)
d: int = foo * "2"         # Pyright: int (正确，不再报错)

print(f"foo + 1 的结果: {c}, 类型: {type(c)}") # 假设Foo的__add__实现返回字符串
print(f"foo * '2' 的结果: {d}, 类型: {type(d)}") # 假设Foo的__mul__实现返回整数