Python Operator Overloading

Python Operator Overloading

In python, some operators behave differently with various types when using python operators with built-in classes. For example, the + operator performs the arithmetic addition on two numbers, joins two strings, or merges two lists.

The feature allowing the same operator to have different meaning according to the context is called operator overloading.

So let us see what will happen when we use those operators with a user-defined class. In the following example, we will try to simulate a point in a 3D coordinate system.

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

p1 = Point(2, 3, 4)
p2 = Point(3, 1, 2)
print(p1 + p2)

Output

 Traceback (most recent call last)
----> 9 print(p1 + p2)

TypeError: unsupported operand type(s) for +: 'Point' and 'Point'

As we can see above, a TypeError was raised because Python did not know how to add two Point together.

To overcome this issue, we can use operator overloading. But first, let us get some notions about special functions.

Python Special Functions

In Python, special functions are class functions that begin with a double underscore __.

The special functions are not the standard functions that we define for a class. For example, when we define the __init__() function, this function is called every time we create a new object of that class.

There are other special functions in Python. To learn more about them, you can visit Python Special Functions.

We can use special functions to make our class compatible with built-in functions.

>>> p1 = Point(1, 3, 2)
>>> print(p1)
<__main__.Point object at 0x7fc6f65c8210>

If we want to print coordinates of the Point object instead of what we got, we can define a special function __str__() in our class that handles how the object gets printed.

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

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

p1 = Point(1, 3, 2)
print(p1)

Output

(1, 3, 2)

As we can see above, now we can output a better presentation of the point.

The __str__() is also invoked when we use the buit-in function str() or format().

>>> str(p1)
(1, 3, 2)

>>> format(p1)
(1, 3, 2)

Here, when str(p1) or format(p1) is called, Python internally calls the p1.__str__() method.

Overloading the + Operator

To overload the + operator in Python, we must implement the __add__() function inside the class. We can write anything inside the __add__ function. But it is recommended to return a Point object of the coordinate sum.

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

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

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

Now let us try to do some addition operation.

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

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

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

p1 = Point(1, 1, 1)
p2 = Point(2, 2, 2)
p3 = Point(3, 3, 3)

print(p1+p2+p3)

Output

(6, 6, 6)

Above, when we use p1 + p2 + p3, Python calls p1.__add__(p2).__add__(p3). And then the addition operation is carried out the way we specified.

We can also overload other operators. The special function that we need to implement is given as follows.

Overloading Comparison Operators

In addition to overloading arithmetic operators, Python offers the possibility to overload comparison operators as well.

Suppose we wanted to implement the greater than symbol > in our Point class.

We will compare the magnitude of these points from the origin and return the result for this purpose. It can be implemented as follows.

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

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

    def __gt__(self, other):
        self_mag = (self.x ** 2) + (self.y ** 2) + (self.z ** 2)
        other_mag = (other.x ** 2) + (other.y ** 2) + (other.z ** 2)
        return self_mag > other_mag

p1 = Point(1, 1, 1)
p2 = Point(2, 2, 2)
p3 = Point(3, 3, 3)

print(p1>p2)
print(p3>p2)
print(p2>p1)

Output

False 
True 
True

We can also overload other comparison operators by implementing special functions as follows.