CODEWITHSUNDEEP
pythonclassoperator-overloading
cdecompilador
cdecompilador

Reputation: 190

Can I use an overloaded operator in the class declaration?

I'm trying to create a simple Vector class in python and I realised that i can't use a previously overloaded operator in a method like in c++. Im having any mistake? or its how python works?

class Vector3D:
    #constructor
    def __init__(self,x,y,z):
        self.x = x
        self.y = y
        self.z = z
    #overloading division operator for just numbers
    def __div__(self,other):
        if type(other) == int or type(other) == float:
            return Vector3D(self.x/other,self.y/other,self.z/other)
        print("Error")
        return
    #Used in the normalize method
    def Magnitude(self):
        return math.sqrt(self.x**2+ self.y**2 + self.z**2)
    # ==> using the division operator in normalization
    def Normalize(self):
        return self / self.Magnitude()
    # <== Throws: TypeError: unsupported operand type(s) for /: 'Vector3D' and 'float'

Upvotes: 0

Views: 111

Answers (1)

jirassimok
jirassimok

Reputation: 4292

In Python 3, use __truediv__ for / and __floordiv__ for //.

This also works in Python 2 with from __future__ import division (which makes division behave as in Python 3).

Here are two additional pieces of advice:

First, rather than checking for the exact types float and int, you can use the classes in the numbers module to check for numeric types.

Second, you should return NotImplemented from an overridden binary operator if it doesn't know what to do with the given value. Then, Python will check if the other operand implements the flipped version of the operator (i.e. __r[name]__) and try that. If that also returns NotImplemented, Python will raise a TypeError. This allows you to make classes that can be used on both sides of the operators.

import numbers

class Vector1D:
  def __init__(self, x):
    self.x = x
  def __repr__(self, x):
    return "Vector1D({})".format(self.x)

  def __mul__(self, other):
    if isinstance(other, numbers.Number):
      return self.__class__(self.x * other)
    return NotImplemented

  def __rmul__(self, other):
    # x * vec1d == vec1d * x
    return self.__mul__(other)

  def __rtruediv__(self, other):
    # You can't divide things by a Vector1D
    return NotImplemented

vec2 = Vector1D(2)

vec2 * 3
# Runs vec2.__mul__(2), gets Vector1D(6)

3 * vec2
# First tries int.__mul__(3, vec2), gets NotImplemented
# Then tries Vector1D.__rmul__(vec2, 3), gets Vector1D(6)

3 / vec2
# First tries int.__truediv__(3, vec2), gets NotImplemented
# Then tries Vector1D.__rtruediv__(vec2, 3), gets NotImplemented
# Python raises TypeError

Upvotes: 2

Related Questions