如何在python中返回不同类型的对象重载sum函数？

Question

我有一个类Data，我想重载__add__函数，并根据我求和的对象的类型来获取不同类型的对象。

玩具示例代码

Data类如下：

class Data(object):

    def __str__(self):
        s = "Data type " + self.type + ": " + self.tag
        return s

    def __init__(self, tag=""):
        self.type = ""
        self.tag = tag
        self.df = pd.Series()

    def createRandomData(self, amountData=10, sizeData=100):
        self.df = pd.DataFrame(np.random.randint(0,sizeData,size=(amountData, 1)))[0]

    def __add__(self, other):
        data = Data()
        data = Data(self.tag + "+" + other.tag)
        data.df = self.df + other.df
        return data

我还有两个不同的对象DataTypeA和DataTypeB，它们继承自Data。

class DataTypeA(Data):

    def __init__(self, tag=""):
        Data.__init__(self, tag)
        self.type = "A"

class DataTypeB(Data):

    def __init__(self, tag=""):
        Data.__init__(self, tag)
        self.type = "B"

预期结果

因此，我想确保当我对相同类型的数据（例如：DataTypeA + DataTypeA）中的两个对象求和时，结果也是该类型的数据。但是，当两个对象属于不同类型（例如：DataTypeA + DataTypeB）时，结果应为Data对象。

总和的结果应该是df中Series的总和，并且标签应该是第一个对象与“ +”串联的标签，最后是第二个标签。

示例

现在，如果我运行以下代码：

a = DataTypeA("data1")
a.createRandomData()
b = DataTypeB("data2")
b.createRandomData()

a2 = a+a
b2 = b+b
c = a+b

print a
print b
print a2
print b2
print c

我得到的结果是：

>>> 
Data type A: data1
Data type B: data2
Data type : data1+data1
Data type : data2+data2
Data type : data1+data2

但是a2（data1 + data1）的类型应为DataTypeA，而不是Data，而b2（data2 + data2）的类型应为DataTypeB而不是Data。

问题

如何重写此代码以实现此目的，这将是最好的方法？我应该改写__add__对象的Data方法还是改写__add__和DataTypeA的重载DataTypeB方法？

Answer 1

一种解决方案是按如下方式重载Data对象__add__方法：

class Data(object):

    def __str__(self):
        s = "Data type " + self.type + ": " + self.tag
        return s

    def __init__(self, tag=""):
        self.type = ""
        self.tag = tag
        self.df = pd.Series()

    def createRandomData(self, amountData=10, sizeData=100):
        self.df = pd.DataFrame(np.random.randint(0,sizeData,size=(amountData, 1)))[0]

    def __add__(self, other):
        data = Data()
        if type(self) == type(other):
            if type(self) == DataTypeA:
                data = DataTypeA(self.tag + "+" + other.tag)
                data.df = self.df + other.df
            elif type(self) == DataTypeB:
                data = DataTypeB(self.tag + "+" + other.tag)
                data.df = self.df + other.df
        else:
            data = Data(self.tag + "+" + other.tag)
            data.df = self.df + other.df
        return data

其余代码将保持不变。

因此，如果您运行问题中提到的相同示例代码，输出将是：

>>> 
Data type A: data1
Data type B: data2
Data type A: data1+data1
Data type B: data2+data2
Data type : data1+data2