Overload arithmetic & comparison operators for element-wise list operations in Python-CodePudding

I have a class as:

class Data:
    def __init__(self, data: dict):
        self.data = data
    
    def __add__(self):

    def __sub__(self):

    def __mul__(self):
   
    def __truediv__(self):

    def __lt__(self):

    def __le__(self):

    def __eq__(self):

    def __ne__(self):

    def __gt__(self):

    def __ge__(self):

I want to perform operations such as

d = {'a': [1, 2, 3], 'b':[4, 5, 6]}
objA = Data(d)
objA.data['a']   2               # this should return element-wise addition=> [3, 4, 5]
objA.data['a']   objA.data['b']  # this should return element-wise addition=> [5, 7, 9]
objA.data['a'] > 2               # this should return element-wise comparison=> [False, False, True]

How can I overload the arithmetic and comparison operators to perform element-wise operations of lists?

I have tried

def __add__(self, l, other):
    return list(map(lambda x: x > other, l)

But get the error: TypeError: can only concatenate list (not "int") to list when doing objA.data['a'] 2

CodePudding user response：

You can use this:

class MathList:
    def __init__(self, data: list):
        self.data = data
    
    def __add__(self, n):
        return [x   n for x in self.data]

    def __sub__(self, n):
        return [x - n for x in self.data]

    def __mul__(self, n):
        return [x * n for x in self.data]

    def __truediv__(self, n):
        return [x / n for x in self.data]


class Data:
    def __init__(self, data: dict):
        self.data = {k: MathList(v) for k, v in data.items()}

objA = Data({
    'a': [1, 2, 3],
    'b': [4, 5, 6]
})

print(objA.data['a']   2) # Output: [3, 4, 5]
print(objA.data['b'] - 2) # Output: [2, 3, 4]
print(objA.data['a'] * 2) # Output: [2, 4, 6]
print(objA.data['b'] / 2) # Output: [2.0, 2.5, 3.0]

This makes a new class called MathList and this performs the operations you want

Edit: you might want to look at numpy, that can also do this

CodePudding user response：

I have defined the data as a list

class Data:
    def __init__(self, data: list):
        self.data = data

    def __add__(self, x):
        self.data.append(x)
        return self

    def __sub__(self, x):
        self.data.remove(x)
        return self

    def __lt__(self, x):
        return [item < x for item in self.data]

    def __le__(self, x):
        return [item <= x for item in self.data]

    def __eq__(self, x):
        return [item == x for item in self.data]

    def __ne__(self, x):
        return [item != x for item in self.data]

    def __gt__(self, x):
        return [item > x for item in self.data]

    def __ge__(self, x):
        return [item >= x for item in self.data]

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


objA = Data([1, 2, 3])
objB = Data([4, 5, 6])
objA = objA   7
print(objA)  # Output: [1, 2, 3, 7]
objA = objA - 7
print(objA)  # Output: [1, 2, 3]
print(objA < 2)  # Output: [True, False, False]
print(objA <= 2)  # Output: [True, True, False]
print(objA == 2)  # Output: [False, True, False]
print(objA != 2)  # Output: [True, False, True]
print(objA < 2)  # Output: [True, False, False]
print(objA >= 2)  # Output: [False, True, True]