28.2 Arithmetic Operators: __add__, __sub__, __mul__, __truediv__
The Role of Arithmetic Dunder Methods
Arithmetic dunder methods are the mechanism by which Python’s data model allows objects to respond to standard operators like +, -, *, and /. When you write a + b, the Python interpreter effectively translates this into a method call: a.__add__(b). This design provides a consistent interface for operator overloading, enabling your custom objects to behave like built-in types. The primary purpose of these methods is not merely to add functionality but to integrate your objects seamlessly into the Python language itself, allowing them to participate in expressions and follow established idioms.
Implementing Basic Arithmetic Operations
To make a custom class support arithmetic operations, you must implement the corresponding dunder methods. These methods should typically return a new instance of the class rather than modifying the current instance (self) in-place. This approach aligns with the expected behavior of arithmetic operations on immutable types (like integers and strings) and avoids surprising side effects.
Consider a Vector class representing a 2D vector:
class Vector:
def __init__(self, x=0, y=0):
self.x = x
self.y = y
def __repr__(self):
return f"Vector({self.x}, {self.y})"
def __add__(self, other):
"""Implements the + operator (self + other)."""
if not isinstance(other, Vector):
return NotImplemented
return Vector(self.x + other.x, self.y + other.y)
def __sub__(self, other):
"""Implements the - operator (self - other)."""
if not isinstance(other, Vector):
return NotImplemented
return Vector(self.x - other.x, self.y - other.y)
def __mul__(self, other):
"""Implements the * operator (self * other)."""
if isinstance(other, (int, float)):
return Vector(self.x * other, self.y * other)
return NotImplemented
def __truediv__(self, other):
"""Implements the / operator (self / other)."""
if isinstance(other, (int, float)):
return Vector(self.x / other, self.y / other)
return NotImplemented
# Example usage:
v1 = Vector(2, 5)
v2 = Vector(3, 1)
print(v1 + v2) # Output: Vector(5, 6)
print(v1 - v2) # Output: Vector(-1, 4)
print(v1 * 3) # Output: Vector(6, 15)
print(v2 / 2) # Output: Vector(1.5, 0.5)
The Importance of Returning NotImplemented
A critical best practice is to return the built-in singleton NotImplemented when an operation is not supported for a given type, as seen in the __mul__ and __truediv__ methods above. This is not the same as raising a NotImplementedError. Returning NotImplemented tells the Python interpreter that the operation is not implemented for these types, which triggers the interpreter to try the reverse operation. For example, if you write 3 * vector_obj, the interpreter first calls int.__mul__(3, vector_obj). When that returns NotImplemented, it then tries vector_obj.__rmul__(3). This mechanism allows for flexible and commutative operator support.
Supporting Reverse Operations (Reflected Methods)
If you try to execute 3 * v1 with the current Vector class, it will fail because the integer 3 doesn’t know how to multiply itself by a Vector. To handle this case, you must implement the reflected (or “right-hand”) methods, such as __rmul__.
class Vector(Vector): # Extending the previous class
def __rmul__(self, other):
"""Implements the reverse * operator (other * self)."""
# Multiplication is commutative, so we can delegate to __mul__
return self.__mul__(other)
# Now this works:
v1 = Vector(2, 5)
print(3 * v1) # Output: Vector(6, 15)
The __rmul__ method is only called if the left operand’s __mul__ returns NotImplemented. This design ensures that the reflected method is a fallback, not the primary path.
Common Pitfalls and Best Practices
- Type Checking and
NotImplemented: Always check the type ofotherand returnNotImplementedfor unsupported types. This is crucial for enabling the reverse method fallback mechanism. - Avoid In-Place Modification: Arithmetic operations are generally expected to be non-mutating. Your
__add__method should create and return a new instance rather than changingself.xandself.ydirectly. Mutating during+would be highly unexpected and a major source of bugs. - Handling Division by Zero: The
__truediv__method must account for the possibility of division by zero. Failing to do so will let a built-inZeroDivisionErrorpropagate, which may be acceptable, but you should be aware of it.def __truediv__(self, other): if isinstance(other, (int, float)): if other == 0: raise ZeroDivisionError("division by zero") return Vector(self.x / other, self.y / other) return NotImplemented - Commutativity and Associativity: The data model does not enforce these mathematical properties. It is the developer’s responsibility to implement the methods in a way that respects them if required. For instance, ensuring
a + b == b + a(commutativity) for your objects is up to your implementation of__add__and__radd__.
In-Place Operators: A Note on __iadd__ and Friends
It’s important to distinguish between + and +=. The latter is an in-place operation. If you do not implement an in-place method like __iadd__, Python will fall back to using __add__ and then assigning the result back to the variable. For mutable objects, this is often inefficient. For a mutable Vector, you might implement __iadd__ to avoid creating a new object.
class Vector(Vector):
def __iadd__(self, other):
"""Implements the += operator."""
if isinstance(other, Vector):
self.x += other.x
self.y += other.y
return self # Important: return self for in-place modification
return NotImplemented
v1 = Vector(2, 5)
v2 = Vector(3, 1)
v1 += v2
print(v1) # Output: Vector(5, 6); v1 was modified in-place.
For immutable objects, relying on the __add__ fallback for += is perfectly acceptable and standard practice.