Operator Overloading with Dunders

Teach your own classes to respond to +, ==, <, [] and more by defining special methods.

Every operator in Python is really a call to a dunder (double-underscore) method. When you write a + b, Python calls a.__add__(b). Define that method on your class and your objects gain the + operator — this is operator overloading.

You writePython calls
a + ba.__add__(b)
a == ba.__eq__(b)
a < ba.__lt__(b)
a[i]a.__getitem__(i)
len(a)a.__len__()
str(a)a.__str__()
class Money:
    def __init__(self, cents):
        self.cents = cents
    def __add__(self, other):
        return Money(self.cents + other.cents)
    def __repr__(self):
        return f'${self.cents/100:.2f}'

print(Money(150) + Money(75))   # $2.25

Do it thoughtfully

Overloading is powerful but easy to abuse. Only define an operator when its meaning is obvious — + for combining vectors, yes; + for something surprising, no. And if you define __eq__, define __hash__ too (or set it to None) so your object behaves correctly in sets and dict keys.

Example

Example · python
from functools import total_ordering

@total_ordering                      # generates <=, >, >= from __eq__ and __lt__
class Vector:
    def __init__(self, x, y):
        self.x, self.y = x, y

    def __add__(self, other):
        if not isinstance(other, Vector):
            return NotImplemented    # let Python try the other side
        return Vector(self.x + other.x, self.y + other.y)

    def __mul__(self, scalar):        # scale by a number
        return Vector(self.x * scalar, self.y * scalar)

    def __abs__(self):                # magnitude, via abs()
        return (self.x ** 2 + self.y ** 2) ** 0.5

    def __eq__(self, other):
        return (self.x, self.y) == (other.x, other.y)

    def __lt__(self, other):          # order by magnitude
        return abs(self) < abs(other)

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

a, b = Vector(1, 2), Vector(3, 4)
print('a + b     :', a + b)          # Vector(4, 6)
print('b * 2     :', b * 2)          # Vector(6, 8)
print('|b|       :', abs(b))         # 5.0
print('a < b     :', a < b)          # True  (sqrt(5) < 5)
print('a >= b    :', a >= b)         # False (free from total_ordering)
print('sorted    :', sorted([b, a])) # [Vector(1, 2), Vector(3, 4)]

# Output:
# a + b     : Vector(4, 6)
# b * 2     : Vector(6, 8)
# |b|       : 5.0
# a < b     : True
# a >= b    : False
# sorted    : [Vector(1, 2), Vector(3, 4)]

When to use it

  • A Money class overloads + and - so currency arithmetic uses natural operator syntax.
  • A Matrix class overloads * so matrix multiplication can be written as A * B.
  • A custom collection class overloads [] (getitem/setitem) to implement a sparse array backed by a dict.

More examples

Arithmetic dunder methods

Overloads + via __add__ so two Money objects can be summed with natural syntax.

Example · python
class Money:
    def __init__(self, amount, currency='USD'):
        self.amount = amount
        self.currency = currency

    def __add__(self, other):
        return Money(self.amount + other.amount, self.currency)

    def __repr__(self):
        return f'{self.currency} {self.amount:.2f}'

a = Money(10.00)
b = Money(5.50)
print(a + b)   # USD 15.50

Comparison dunders

Defines __eq__ and __lt__ and uses @total_ordering to auto-generate the remaining comparisons.

Example · python
from functools import total_ordering

@total_ordering
class Version:
    def __init__(self, major, minor):
        self.major = major
        self.minor = minor

    def __eq__(self, other):
        return (self.major, self.minor) == (other.major, other.minor)

    def __lt__(self, other):
        return (self.major, self.minor) < (other.major, other.minor)

v1 = Version(1, 9)
v2 = Version(2, 0)
print(v1 < v2)   # True
print(v1 > v2)   # False

Container dunders: getitem and setitem

Implements [] read/write access via __getitem__ and __setitem__ backed by a dict.

Example · python
class SparseArray:
    def __init__(self):
        self._data = {}

    def __setitem__(self, index, value):
        self._data[index] = value

    def __getitem__(self, index):
        return self._data.get(index, 0)

arr = SparseArray()
arr[100] = 42
print(arr[100])   # 42
print(arr[200])   # 0  (default)

Discussion

  • Be the first to comment on this lesson.