Динамически добавлять поддержку магической арифметической функции в классе Python

Question

Допустим, у меня есть класс Python A:

class A:
    def __init__(self, matrix, metadata: list):
        self.matrix = np.array(matrix)
        self.metadata = metadata
    #... 

Теперь я хочу, чтобы все арифметические операции работали для моего класса. Они должны были просто перевести операцию на matrix, то есть так:

    def __add__(self, other):
        if isinstance(other, type(self)):
            raise ValueError("Not allowed.")
        else:
            return A(
                matrix=self.matrix.__add__(other),
                metadata=self.metadata,
            )

Теперь проблема в том, что мне приходится повторять практически одинаковый код для каждой магической арифметической функции, т.е. __add__, __sub__, __mul__, __truediv__, __pow__, __radd__, __rsub__, __rmul__, __rtruediv__, __iadd__, __isub__, __imul__, __itruediv__, __abs__, __round__, __floor__, __ceil__, __trunc__. Что приводит к большому количеству повторяющихся кодов.

Как я могу определить их динамически в цикле for? как

magic_functions = ["__add__", "__sub__", ...]
for magic_function in magic_functions:
    # define the method and add it to the class

Answer 1

Я хотел бы предложить вам использовать decorator в этой ситуации.Может быть, это не так коротко, но вы сохраните читабельность вашего кода.

import numpy as np

def decorator(fn):
    def ret_fn(*args, **kwargs):
        if isinstance(args[1], type(args[0])):
            raise ValueError("Not allowed.")
        else:
            return fn(*args, **kwargs)

    return ret_fn

class A:
    def __init__(self, matrix, metadata: list):
        self.matrix = np.array(matrix)
        self.metadata = metadata

    @decorator
    def __add__(self, other):
        return A(
            matrix=self.matrix.__add__(othe),
            metadata=self.metadata,
        )

Результат:

>>> a1 = A([[1], [2]], [])
>>> a2 = a1 + [[3], [4]]
>>> print(a2.matrix)
[[4]
 [6]]
>>> a1 + a1
Traceback (most recent call last):
...
    raise ValueError("Not allowed.")
ValueError: Not allowed.

Я не знаю, сколько различий между вашими функциями, ноВы можете переписать декоратор и функцию довольно минималистично:

декоратор

def decorator(fn):
    def ret_fn(*args, **kwargs):
        if isinstance(args[1], type(args[0])):
            raise ValueError("Not allowed.")
        else:
            return A(
                    matrix=fn(*args, **kwargs),
                    metadata=args[0].metadata,
                )

    return ret_fn

метод

    @decorator
    def __add__(self, other):
        return self.matrix.__add__(other)

Answer 2

Эта (широкая) проблема является целью модуля operator:

import operator
def mkop(f):    # the usual scope for a closure
  def op(self,o):
    if isinstance(o,type(self)): raise …
    return type(self)(matrix=f(self.matrix,o),
                      metadata=self.metadata)
  return op
for op in ['add','sub',…]:
  setattr(A,"__%s__"%op,mkop(getattr(operator,op)))

Вы также можете использовать locals()[…]=mkop(…) (в одном из его редких безопасных применений), чтобы выполнить вышеизложенное при определении класса.