As an example, assume a codebase includes the following class. Instances represent complex numbers that can be multiplied:
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
Real numbers are a special-case of complex numbers, so the codebase also includes a subclass for them:
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
I believe the above is all idiomatic Python, and it all works fine. Now assume that I want to add a class for imaginary numbers to the existing code. Imaginary
is a subclass of Complex
and a sibling class of 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
Unfortunately one of these assertions fails. type(Real(0) * Imaginary(0))
is Complex
rather than Imaginary
because Python performs the multiplication by calling Real.__mul__
instead of Imaginary.__rmul__
.
Is there any way to write Imaginary
so that the overloaded operators all work as expected, without modifying Complex
or Real
?
Or should Complex
and Real
be written differently to enable this to work (without requiring them to have knowledge of Imaginary
)?
from How can I make operator overloading work correctly when adding a sibling class?
No comments:
Post a Comment