Добавление двух отдельных подклассов в качестве их общего суперкласса

Question

У меня есть класс с большим количеством подклассов в моем коде. Рассмотрим следующий код:

class Dataset:
  def __init__(self, samples):
    self.data = samples

  def __add__(self, other):
    if isinstance(other, type(self)):
      return type(self)(self.data + other.data)
    return NotImplemented

  def __radd__(self, other):
    if isinstance(other, type(self)):
      return type(self)(other.data + self.data)
    return NotImplemented

  def __str__(self):
    return str(self.data)


class DatasetA(Dataset):
  def __init__(self, samples):
    super().__init__(samples)
    self.a = len(self.data)

  def __str__(self):
    return f"{self.a}: {self.data}"


class DatasetB(Dataset):
  def __init__(self, samples):
    super().__init__(samples)
    self.b = sum(self.data)

  def __str__(self):
    return f"{self.data} (sum: {self.b})"


d = Dataset([1, 2, 3])
a = DatasetA([4, 5])
b = DatasetB([6, 7])

print(d + d)
print(d + a)
print(a + d)
print(d + b)
print(b + d)
print(a + a)
print(b + b)
print(a + b)
print(b + a)

Идея состоит в том, чтобы определить метод __add__ в суперклассе, который не требует переопределения в каждом из подклассов и все равно будет корректно добавлять их (например, два DatasetB s должно складываться с другим DatasetB). Это работает правильно (первые 7 print с подойдут), однако я хотел бы реализовать одну дополнительную функциональность, представленную в последних 2 print с.

Я бы хотел, чтобы любые два разных подклассы, чтобы сложить в первый общий суперкласс. Например, a + b должен привести к экземпляру Dataset. Также, если мы добавим экземпляры двух подклассов DatasetB, результатом должен быть экземпляр DatasetB.

Я попытался изменить return NotImplemented на return super().__add__(other) (и аналогично для __radd__) , но это уже привело к AttributeError для 3-го print оператора.

Есть ли способ реализовать эту желаемую функциональность, не нарушая существующую (то есть, все еще правильно выполняются первые 7 отпечатков), и без необходимости явно переопределять __add__ в каждом из подклассов?

Answer 1

Вы можете изменить:

return type(self)(self.data + other.data)

на:

klass = next(c for c in type(self).mro() if c in type(other).mro())
return klass(self.data + other.data)

Это выберет наиболее специфицированный c суперкласс, общий для двух, используя разрешение метода заказ .