[Classes] Class Intermediate: Operator Overloading

[ichabod801]

Okay, then what operations do we normally do on vectors? We normally add them together and multiply them by scalars. Just as the comparison operators have methods allowing you to override them, so do addition, multiplication, and so on:

__add__ overrides +
__sub__ overrides -
__mul__ overrides *
__truediv__ overrides /
__floordiv__ overrides //
__mod__ overrides %
__pow__ overrides **
__lshift__ overrides <<
__rshift__ overrides >>
__and__ overrides &
__xor__ overrides ^
__or__ overrides |

Note that __truediv__ is not supported in Python 2.x unless you have imported division from future. In 2.x, use __div__ for /, but realize that __div__ doesn't work in Python 3.x. If you are trying to develop for both, implement __div__ and __truediv__.

But we're just going to implement __add__ and __mul__ for our Vector class:

   def __add__(self, other):
      """
      Addition.
      
      Parameters:
      other: the other summand. (Vector, list, or tuple)
      """
      if isinstance(other, Vector):
         return Vector(self.x + other.x, self.y + other.y)
      elif isinstance(other, (tuple, list)):
         return Vector(self.x + other[0], self.y + other[1])
      else:
         return NotImplemented
   
   def __mul__(self, other):
      """
      Scalar multiplication.
      
      Parameters:
      other: The multiplier (float or int)
      """
      # vector * number
      if isinstance(other, (float, int)):
         return Vector(other * self.x, other * self.y)
      # vector * other
      else:
         return NotImplemented

When we try it out we get:

>>> a = Vector(8, 1)
>>> b = Vector(3, 4)
>>> a + b
Vector(11, 5)
>>> b + a
Vector(11, 5)
>>> a + b == b + a
True
>>> a + (3, 4)
Vector(11, 5)
>>> 5 * a
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unsupported operand type(s) for *: 'int' and 'Vector'

So why didn't the multiplication work? Here's a clue:

>>> (3, 4) + a
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: can only concatenate tuple (not "Vector") to tuple
>>> a * 5
Vector(40, 5)

So __add__ and __mul__ are only working when the vector is on the left side. Here's what is going on: When Python sees 'spam + eggs', it tries 'spam.__add__(eggs)'. If that returns NotImplemented, it tries 'eggs.__radd__(spam)' (radd for right addition). Now, if all you are ever going to do is add Vectors to Vectors, you just need __add__. But when you start mixing in tuples and lists you have to accont for __radd__ as well. If you were always going to have the scalar on the left, as is standard in mathematics, all you would ever need is __rmul__. But then you'd run the risk of someone typing in 'a * 5', so it's best to implement both.

However, it's easy to do so:

   def __radd__(self, other):
      """
      Right addition.
      
      Parameters:
      other: The other summand (Vector, tuple, or list)
      """
      return self.__add__(other)
         
   def __rmul__(self, other):
      """
      Right multiplication.
      
      Parameters:
      other: The other multiplier (float or int)
      """
      return self.__mul__(other)

There may be objects for which a + b != b + a, but that is not the case with Vectors. So we can just use the functionality we already wrote for __add__ and __mul__ to handle the right hand cases. Again, that means that if we want to change that functionality, we only need to change it in one place. Now things work as expected:

>>> a = Vector(8, 1)
>>> (3, 4) + a
Vector(11, 5)
>>> 5 * a
Vector(40, 5)
>>> 3 * a == a + a + a
True
>>> a * (3, 4)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: can't multiply sequence by non-int of type 'Vector'

We get the last TypeError because we didn't implement multiplication by tuples, lists, or Vectors. But that's the same TypeError the other errors were causing, so we are being Pythonic by causing the expected error.