CODEWITHSUNDEEP
pythonoperator-overloading
boymeetscode
boymeetscode

Reputation: 825

How can I overload operators so that type on the left/right does not matter?

A simple example - I want to have a point class that describes a point in 2 dimensions. And I want to be able to add two points together... as well as multiply two points together (don't ask me why), or multiply a point by a scalar. For right now, I'll only implement it as if the scalar is an integer but fraction or floats is also trivial to do.

class Point:
    def __init__(self, x=0, y=0):
        self.x = x
        self.y = y

    def __str__(self):
        return "({0},{1})".format(self.x, self.y)

    def __add__(self, other):
        x = self.x + other.x
        y = self.y + other.y
        return Point(x, y)

    def __mul__(self, other):
        if isinstance(other, Point):
            x = self.x * other.x
            y = self.y * other.y
            return Point(x, y)
        elif isinstance(other, int):
            x = self.x * other
            y = self.y * other
            return Point(x, y)

So, this works when I perform:

>>> p1 = Point(2, 3)
>>> p2 = Point(-1, 2)
>>> print(p1*p2)
(-2,6)
>>>print(p1*4)
(8,12)

but it does NOT work when I reverse the order of scalar and Point object:

>>>print(4*p1)
Traceback (most recent call last):   
  File "<input>", line 1, in <module> TypeError: unsupported operand type(s) for *:
'int' and 'Point'

How can I write the code where it doesn't matter if I write '4 * p1' or 'p1 * 4' that I'll still execute the same code and return the same answer? Do I accomplish this by overloading the mul operator for int objects or is there another way?

Note: code for my short example was borrowed from https://www.programiz.com/python-programming/operator-overloading

Upvotes: 7

Views: 1866

Answers (1)

boymeetscode
boymeetscode

Reputation: 825

(As I was about to submit the question, I was tagging, and found the answer. I think it is worthwhile to document it here so others can easily find it.)

Define __rmul__(self, other). This stands for right-multiply. When the object on the left fails to multiply (in the example above the integer doesn't know how to multiply the Point class on the right) Python will look at the object on the right to see if the __rmul__(self, other) special method is defined, and does it work? If it does, it will use this implementation instead.

For classes that are commutative (i.e. you can multiply either AB or BA and get the same results) you can define it as:

def __mul__(self, other):
    if isinstance(other, Point):
        x = self.x * other.x
        y = self.y * other.y
        return Point(x, y)
    elif isinstance(other, int):
        x = self.x * other
        y = self.y * other
        return Point(x, y)

def __rmul__(self, other):
    return self.__mul__(other)

Upvotes: 9

Related Questions