Python Operator Overloading

Summary: in this tutorial, you’ll learn Python operator overloading and how to use it to make your objects work with built-in operators.

Introduction to the Python operator overloading

Suppose you have a 2D point class with x and y coordinate attributes:

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

    def __str__(self):
        return f'({self.x},{self.y})'Code language: Python (python)

To add two Point2D objects, you can define an add() method as follows:

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

    def __str__(self):
        return f'({self.x},{self.y})'

    def add(self, point):
        if not isinstance(point, Point2D):
            raise ValueError('The other must be an instance of the Point2D')

        return Point2D(self.x + point.x, self.y + point.y)Code language: Python (python)

The add() method raises an error if the point is not an instance of the Point2D class. Otherwise, it returns a new Point2D object whose x and y coordinates are the sums of x and y coordinates of two points.

The following creates two instances of the Point2D class and use the add() method to add two points:

a = Point2D(10, 20)
b = Point2D(15, 25)
c = a.add(b)

print(c)Code language: Python (python)

Output:

(25,45)Code language: Python (python)

This code works perfectly fine. But Python has a better way to implement it. Instead of using the add() method, you can use the built-in operator (+) like this:

c = a + bCode language: Python (python)

When you use the + operator on the Point2D object, Python will call the special method __add__() on the object. The following calls are equivalent:

c = a + b
c = a.__add__(b)Code language: Python (python)

The __add__() method must return a new instance of the Point2D object.

The ability to use the built-in operator (+) on a custom type is known as operator overloading.

The following shows the Point2D class that implements the __add__() special operator to support the + operator:

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

    def __str__(self):
        return f'({self.x},{self.y})'

    def __add__(self, point):
        if not isinstance(point, Point2D):
            raise ValueError('The other must be an instance of the Point2D')

        return Point2D(self.x + point.x, self.y + point.y)


if __name__ == '__main__':
    a = Point2D(10, 20)
    b = Point2D(15, 25)
    c = a + b
    print(c)Code language: Python (python)

Output:

(25,45)

Special methods for operator overloading

The following shows the operators with their corresponding special methods:

For example, you can implement the __sub__() method in the Point2D to support subtraction (-) of two points:

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

    def __str__(self):
        return f'({self.x},{self.y})'

    def __add__(self, point):
        if not isinstance(point, Point2D):
            raise ValueError('The other must be an instance of the Point2D')

        return Point2D(self.x + point.x, self.y + point.y)

    def __sub__(self, other):
        if not isinstance(other, Point2D):
            raise ValueError('The other must be an instance of the Point2D')

        return Point2D(self.x - other.x, self.y - other.y)


if __name__ == '__main__':
    a = Point2D(10, 20)
    b = Point2D(15, 25)
    c = b - a
    print(c)Code language: Python (python)

Overloading inplace opeators

Some operators have the inplace version. For example, the inplace version of + is +=.

For the immutable type like a tuple, a string, a number, the inplace operators perform calculations and don’t assign the result back to the input object.

For the mutable type, the inplace operator performs the updates on the original objects directly. The assignment is not necessary.

Python also provides you with a list of special methods that allows you to overload the inplace operator:

Let’s take an example of overloading the += operator.

Suppose you have a cart object and you want to add an item to the cart. To do you, you can define an add() method to the Cart class and use it like this:

cart.add(item)Code language: Python (python)

Alternatively, you can implement the += operator in the Cart class. It allows you to add an item to the cart as follows:

cart += itemCode language: Python (python)

To support the += operator, you need to implement the __iadd__ special method in the Cart class.

First, define the Item class that has three attributes name, quantity, and price. Also, it has an amount property that returns the subtotal of the item:

class Item:
    def __init__(self, name, qty, price):
        self.name = name
        self.qty = qty
        self.price = price

    @property
    def amount(self):
        return self.qty * self.price

    def __str__(self):
        return f'{self.name} {self.qty} ${self.price} ${self.amount}'Code language: Python (python)

Second, define the Cart class that implements the __iadd__ method:

class Cart:
    def __init__(self):
        self.items = []

    def __iadd__(self, item):
        if not isinstance(item, Item):
            raise ValueError('The item must be an instance of Item')

        self.items.append(item)
        return self

    @property
    def total(self):
        return sum([item.amount for item in self.items])

    def __str__(self):
        if not self.items:
            return 'The cart is empty'

        return '\n'.join([str(item) for item in self.items])Code language: Python (python)

In the __iadd__ method, we raise a ValueError if the item is not an instance of the Item class. Otherwise, we add the item to the items list attribute.

The total property returns the sum of all items.

The __str__ method returns the string 'The cart is empty' if the cart has no item. Otherwise, it returns a string that contains all items separated by a newline.

Third, use the += operator to add an item to the cart:

if __name__ == '__main__':
    cart = Cart()

    cart += Item('Apple', 5, 2)
    cart += Item('Banana', 20, 1)
    cart += Item('Orange', 10, 1.5)

    print(cart)
    # print the total line
    print('-' * 30)
    print('Total: $', cart.total)Code language: Python (python)

Output:

Apple   5       $2      $10
Banana  20      $1      $20
Orange  10      $1.5    $15.0
------------------------------
Total: $ 45.0Code language: Python (python)

Summary

• Opeartor overloading allows a class to use built-in operators.