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 write | Python calls |
|---|---|
a + b | a.__add__(b) |
a == b | a.__eq__(b) |
a < b | a.__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.25Do 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
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.
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.50Comparison dunders
Defines __eq__ and __lt__ and uses @total_ordering to auto-generate the remaining comparisons.
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) # FalseContainer dunders: getitem and setitem
Implements [] read/write access via __getitem__ and __setitem__ backed by a dict.
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