简体   繁体   中英

Dynamically add support for arithmetic magic function in a python class

Let say I have a python class A :

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

Now I want all arithmetic operations work for my class. They supposed to simply translate the operation to matrix , ie like so:

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

Now the problem is that I have to repeate almost the same code for each arithmetic magic function, ie __add__, __sub__, __mul__, __truediv__, __pow__, __radd__, __rsub__, __rmul__, __rtruediv__, __iadd__, __isub__, __imul__, __itruediv__, __abs__, __round__, __floor__, __ceil__, __trunc__ . Which leads to a lot of repeating code.

How can I define them dynamically in a for loop? like

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

This (broad) sort of problem is the purpose of the operator module:

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)))

You can also use locals()[…]=mkop(…) (in one of its rare safe uses) to do the above while defining the class.

I would like to suggest to you use the decorator in this situation. May be this is not so short but you will save the readability of your code.

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,
        )

The result:

>>> 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.

I don't know how many difference between your functions but you may rewrite decorator and the function pretty minimalistic:

decorator

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

the method

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

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM