Supporting Operator Overloading in Custom Classes

First up today, you’ll explore the special methods you need to perform operations using Python’s operators on custom classes.

This text is part of a Real Python tutorial by Leodanis Pozo Ramos.

Supporting Operator Overloading in Custom Classes

Python has multiple types of operators, which are special symbols, combinations of symbols, or keywords that designate some type of computation. Internally, Python supports operators with special methods. For example, the .__add__() special method supports the plus operator (+) as you already saw.

In practice, you’ll take advantage of these methods that are behind operators for something that’s known as operator overloading.

Operator overloading means providing additional functionality to the operators. You can do this with most built-in types and their specific supported operators. However, that’s not all you can do with the special methods that support Python operators. You can also use these methods to support some operators in your custom classes.

In the following sections, you’ll learn about the specific special methods that support Python operators, including arithmetic, comparison, membership, bitwise, and augmented operators. To kick things off, you’ll start with arithmetic operators, which are arguably the most commonly used operators.

Arithmetic Operators

Arithmetic operators are those that you use to perform arithmetic operations on numeric values. In most cases, they come from math, and therefore, Python represents them with the usual math signs.

In the following table, you’ll find a list of the arithmetic operators in Python and their supporting magic methods:

Note that all these methods take a second argument called other. In most cases, this argument should be the same type as self or a compatible type. If that’s not the case, then you can get an error.

To illustrate how you can overload some of these operators, get back to the Storage class. Say that you want to make sure that when you add or subtract two instances of this class, both operands have the same unit:

class Storage(float):
    def __new__(cls, value, unit):
        instance = super().__new__(cls, value)
        instance.unit = unit
        return instance

    def __add__(self, other):
        if not isinstance(other, type(self)):
            raise TypeError(
                "unsupported operand for +: "
                f"'{type(self).__name__}' and '{type(other).__name__}'"
            )
        if not self.unit == other.unit:
            raise TypeError(
                f"incompatible units: '{self.unit}' and '{other.unit}'"
            )

        return type(self)(super().__add__(other), self.unit)

In the .__add__() method, you first check if other is also an instance of Storage. If not, then you raise a TypeError exception with an appropriate error message. Next, you check if both objects have the same unit. If not, then you raise a TypeError again. If both checks pass, then you return a new instance of Storage that you create by adding the two values and attaching the unit.

Here’s how this class works:

>>> from storage import Storage

>>> disk_1 = Storage(500, "GB")
>>> disk_2 = Storage(1000, "GB")
>>> disk_3 = Storage(1, "TB")

>>> disk_1 + disk_2
1500.0

>>> disk_2 + disk_3
Traceback (most recent call last):
    ...
TypeError: incompatible units: 'GB' and 'TB'

>>> disk_1 + 100
Traceback (most recent call last):
    ...
TypeError: unsupported operand for +: 'Storage' and 'int'

In this example, when you add two objects that have the same unit, you get the correct value. If you add two objects with different units, then you get a TypeError telling you that the units are incompatible. Finally, when you try to add a Storage instance with a built-in numeric value, you get an error as well because the built-in type isn’t a Storage instance.

As an exercise, would you like to implement the .__sub__() method in your Storage class? Check the collapsible section below for a possible implementation:

class Storage(float):
    def __new__(cls, value, unit):
        instance = super().__new__(cls, value)
        instance.unit = unit
        return instance

    def __add__(self, other):
        if not isinstance(other, type(self)):
            raise TypeError(
                "unsupported operand for +: "
                f"'{type(self).__name__}' and '{type(other).__name__}'"
            )
        if not self.unit == other.unit:
            raise TypeError(
                f"incompatible units: '{self.unit}' and '{other.unit}'"
            )

        return type(self)(super().__add__(other), self.unit)

    def __sub__(self, other):
        if not isinstance(other, type(self)):
            raise TypeError(
                "unsupported operand for -: "
                f"'{type(self).__name__}' and '{type(other).__name__}'"
            )
        if not self.unit == other.unit:
            raise TypeError(
                f"incompatible units: '{self.unit}' and '{other.unit}'"
            )

        return type(self)(super().__sub__(other), self.unit)

In the above example, you overloaded the addition operator (+) on the built-in float type. You can also use the .__add__() method and any other operator method to support the corresponding operator in a custom class.

Consider the following example of a Distance class: