添加同级类时，如何使运算符重载正常工作？

Question

例如，假设代码库包含以下类。实例代表可以相乘的复数：

class Complex:
    def __init__(self, re, im):
        self.re, self.im = re, im
    def __mul__(lhs, rhs):
        if isinstance(rhs, Complex): return Complex(...)
        return NotImplemented

assert type(Complex(0, 0) * Complex(0, 0)) is Complex

实数是复数的特例，因此代码库还包含了它们的子类：

class Real(Complex):
    def __init__(self, re):
        self.re, self.im = re, 0
    def __mul__(lhs, rhs):
        if isinstance(rhs, Real): return Real(...)
        if isinstance(rhs, Complex): return Complex(...)
        return NotImplemented
    def __rmul__(rhs, lhs):
        if isinstance(lhs, Complex): return Complex(...)
        return NotImplemented

assert type(Real(0) * Real(0)) is Real
assert type(Real(0) * Complex(0, 0)) is Complex
assert type(Complex(0, 0) * Real(0)) is Complex

我相信以上都是惯用的 Python，而且一切正常。现在假设我想在现有代码中添加一个虚数类。 Imaginary 是Complex 的子类和Real 的兄弟类：

class Imaginary(Complex):
    def __init__(self, im):
        self.re, self.im = 0, im
    def __mul__(lhs, rhs):
        if isinstance(rhs, Imaginary): return Real(...)
        if isinstance(rhs, Real): return Imaginary(...)
        if isinstance(rhs, Complex): return Complex(...)
        return NotImplemented
    def __rmul__(rhs, lhs):
        if isinstance(lhs, Real): return Imaginary(...)
        if isinstance(lhs, Complex): return Complex(...)
        return NotImplemented

assert type(Imaginary(0) * Imaginary(0)) is Real
assert type(Imaginary(0) * Real(0)) is Imaginary
assert type(Imaginary(0) * Complex(0, 0)) is Complex
assert type(Real(0) * Imaginary(0)) is Imaginary  # AssertionError
assert type(Complex(0, 0) * Imaginary(0)) is Complex

不幸的是，其中一个断言失败了。 type(Real(0) * Imaginary(0)) 是 Complex 而不是 Imaginary，因为 Python 通过调用 Real.__mul__ 而不是 Imaginary.__rmul__ 来执行乘法运算。

有没有办法写Imaginary，让重载的操作符都按预期工作，而不需要修改Complex或Real？

或者应该以不同的方式编写 Complex 和 Real 以使其能够正常工作（不需要他们了解 Imaginary）？

Answer 1

Dymtro 是正确的，使用 isinstance 可以代替类型/类检查，但是您不应该直接检查 obj.__class__。相反，请使用 type(obj) is Classname 进行检查，就像您使用 Assert 检查一样！在这里，我用 0 代替实际的复杂数学：

class Complex:
    def __init__(self, re, im):
        self.re, self.im = re, im
    def __mul__(lhs, rhs):
        if type(rhs) is Complex: return Complex(0,0)
        return NotImplemented

assert type(Complex(0, 0) * Complex(0, 0)) is Complex

class Real(Complex):
    def __init__(self, re):
        self.re, self.im = re, 0
    def __mul__(lhs, rhs):
        if type(rhs) is Real: return Real(0)
        if type(rhs) is Complex: return Complex(0,0)
        return NotImplemented
    def __rmul__(rhs, lhs):
        if type(lhs) is Complex: return Complex(0,0)
        return NotImplemented

assert type(Real(0) * Real(0)) is Real
assert type(Real(0) * Complex(0, 0)) is Complex
assert type(Complex(0, 0) * Real(0)) is Complex

class Imaginary(Complex):
    def __init__(self, im):
        self.re, self.im = 0, im
    def __mul__(lhs, rhs):
        if type(rhs) is Imaginary: return Real(0)
        if type(rhs) is Real: return Imaginary(0)
        if type(rhs) is Complex: return Complex(0,0)
        return NotImplemented
    def __rmul__(rhs, lhs):
        if type(lhs) is Real: return Imaginary(0)
        if type(lhs) is Complex: return Complex(0,0)
        return NotImplemented

assert type(Imaginary(0) * Imaginary(0)) is Real
assert type(Imaginary(0) * Real(0)) is Imaginary
assert type(Imaginary(0) * Complex(0, 0)) is Complex
assert type(Real(0) * Imaginary(0)) is Imaginary  # (No longer) AssertionError
assert type(Complex(0, 0) * Imaginary(0)) is Complex

替代方法 1

就像我在上面的评论中所说，Python 将首先尝试使用参数评估 LHS 的 __mul__。只有当 NotImplemented 失败时，它才会尝试 RHS 的__rmul__。见：Under what circumstances are __rmul__ called?

您可以尝试在 Real 不知道具体关于 Imaginary 的情况下先强制另一种方式，但认为可能有另一个对象应该优先于 Complex，如下所示：

class Real(Complex):
    def __init__(self, re):
        self.re, self.im = re, 0
    def __mul__(lhs, rhs):
        if isinstance(rhs, Real): return Real(0)
        try:
            if hasattr(rhs,"__rmul__"):
                return rhs.__rmul__(lhs)
        except NotImplementedError:
            pass
        if isinstance(rhs, Complex): return Complex(0,0)
        return NotImplemented
    def __rmul__(rhs, lhs):
        if isinstance(lhs, Complex): return Complex(0,0)
        return NotImplemented

替代方法 2

如果你反过来做，Complex 继承自 Real 和 Imaginary 的组合，你可以按照你已经完成的方式逐步构建它，第一个类只知道它自己，第二个类知道自己和第一类，最后第三类知道每个人。

class Real:
    def __init__(self, re):
        self.re = re
    def __mul__(lhs, rhs):
        if isinstance(rhs, Real): return Real(lhs.re*rhs.re)
        return NotImplemented

assert type(Real(0) * Real(0)) is Real

class Imaginary:
    def __init__(self, im):
        self.im = im
    def __mul__(lhs, rhs):
        if isinstance(rhs, Imaginary): return Real(lhs.im * rhs.im * -1)
        if isinstance(rhs, Real): return Imaginary(lhs.im * rhs.re)
        return NotImplemented
    def __rmul__(rhs, lhs):
        if isinstance(lhs, Real): return Imaginary(lhs.re * rhs.im)
        return NotImplemented

class Complex(Real,Imaginary):
    def __init__(self, re, im):
        self.re, self.im = re, im
    def __mul__(lhs, rhs):
        if isinstance(rhs, Real): return Complex(lhs.re * rhs.re,lhs.im * rhs.re)
        if isinstance(rhs, Imaginary): return Complex(-1*rhs.im * lhs.im,rhs.im * lhs.re)
        if isinstance(rhs, Complex): return Complex(lhs.re*rhs.re - lhs.im*rhs.im,lhs.re*rhs.im+rhs.re*lhs.im)
        return NotImplemented
    def __rmul__(rhs,lhs):
        if isinstance(lhs, Real): return Complex(lhs.re*rhs.re, lhs.re*rhs.im)
        if isinstance(lhs, Imaginary): return Complex(-1 * lhs.im*rhs.im, lhs.im*rhs.re)

assert type(Complex(0, 0) * Complex(0, 0)) is Complex
assert type(Real(0) * Complex(0, 0)) is Complex
assert type(Complex(0, 0) * Real(0)) is Complex


assert type(Imaginary(0) * Imaginary(0)) is Real
assert type(Imaginary(0) * Real(0)) is Imaginary
assert type(Imaginary(0) * Complex(0, 0)) is Complex
assert type(Real(0) * Imaginary(0)) is Imaginary  # NOT an AssertionError
assert type(Complex(0, 0) * Imaginary(0)) is Complex

Answer 2

isinstance是问题的根源

定义

问题显示为isinstance“返回一个对象是类的实例还是其子类的实例。”

>>> help(isinstance)

isinstance(obj, class_or_tuple, /)
    Return whether an object is an instance of a class or of a subclass thereof.
    
    A tuple, as in ``isinstance(x, (A, B, ...))``, may be given as the target to
    check against. This is equivalent to ``isinstance(x, A) or isinstance(x, B)
    or ...`` etc.

什么意思？

>>> class A:
    pass

>>> class B(A):
    pass

>>> b = B()

>>> isinstance(b, A)
True # `b` instance is an instance of the sublass A

>>> isinstance(b, B)
True # `b` instance is an instance of the class 

为什么assert 失败了？

class Real(Complex):
    def __init__(self, re):
        self.re, self.im = re, 0
    def __mul__(lhs, rhs):
        if isinstance(rhs, Real): return Real(...)
        if isinstance(rhs, Complex): return Complex(...) 
        #                  ^^^^^^^
        # Complex is an instance of Complex
        # Real is an instance of Complex
        # Imagionary is an instance of Complex
        #
        # Any subclass of Complex will return True here.

        return NotImplemented
    def __rmul__(rhs, lhs):
        if isinstance(lhs, Complex): return Complex(...)
        return NotImplemented

解决方案

这里的一个解决方案是比较类而不是isinstance 所以，

if isinstance(rhs, Complex): return Complex(...)

变成

if rhs.__class__ == Complex: return Complex(...)

这里比较的是确切的类，而不是整个层次结构。

传递所有断言的完整代码

class Complex:
    def __init__(self, re, im):
        self.re, self.im = re, im

    def __mul__(lhs, rhs):
        if rhs.__class__ == Complex: return Complex(1, 1)
        return NotImplemented


assert type(Complex(0, 0) * Complex(0, 0)) is Complex


class Real(Complex):
    def __init__(self, re):
        super().__init__(re, 0)

    def __mul__(lhs, rhs):
        if rhs.__class__ == Real: return Real(1)
        if rhs.__class__ == Complex: return Complex(1, 0)
        return NotImplemented

    def __rmul__(rhs, lhs):
        if lhs.__class__ == Complex: return Complex(1, 0)
        return NotImplemented


assert type(Real(0) * Real(0)) is Real
assert type(Real(0) * Complex(0, 0)) is Complex
assert type(Complex(0, 0) * Real(0)) is Complex


class Imaginary(Complex):
    def __init__(self, im):
        super().__init__(0, im)

    def __mul__(lhs, rhs):
        if rhs.__class__ == Imaginary: return Real(1)
        if rhs.__class__ == Real: return Imaginary(1)
        if rhs.__class__ == Complex: return Complex(1, 1)
        return NotImplemented

    def __rmul__(rhs, lhs):
        if lhs.__class__ == Real: return Imaginary(1)
        if lhs.__class__ == Complex: return Complex(1, 1)
        return NotImplemented


assert type(Imaginary(0) * Imaginary(0)) is Real
assert type(Imaginary(0) * Real(0)) is Imaginary
assert type(Imaginary(0) * Complex(0, 0)) is Complex
assert type(Real(0) * Imaginary(0)) is Imaginary
assert type(Complex(0, 0) * Imaginary(0)) is Complex