Arithmetic Methods in Python Classes

by John | November 26, 2020

 

In this article we will discuss and show examples of using arithmetic operations for objects in Python. To make it clear how this could be useful in a real project, we will give a small example at the end of the document to put things into context.

 

If we want to use arithmetic operators on our Python classes, there are a number of built in double underscore methods that we can define to let Python know what we want to achieve. Let's recall some of the basic arithmetic operators everyone should be familiar with. 

 

x = 10
y = 5

#simple addition
z = x + y
#inplace addition
x += 1
# multiplication
x*2
#inplace multiplication
y *= 3

 

You can print the results from the statements above, if you are unsure about the in-place operations. Let's create an example class to show how this can be done with custom classes. 

 

class ArithmeticOps:
    def __init__(self, total):
        self._total = total
        
    def __repr__(self):
        return f"ArithmeticOps(total={self._total})"
        
    

a = ArithmeticOps(10)    

 

Now that we have our object defined let's try to add an integer to the total. Ideally we want to be able to just write instance (a) + value, so let's try that out. 

 

a + 10 

# returns : TypeError: unsupported operand type(s) for +: 'ArithmeticOps' and 'int'

 

Python returns a TypeError when we try to add to our instance. This may seem strange since we only have one number, but consider a case when we have many instance attributes, in this case it is not clear which Python should add the value to. 

 

 

__add__ method

 

If we want to be able to add to our class in the instance + value manner we have tried above we need to implement the __add__ method. Take a look below at how we do this for our example class. 

 

class ArithmeticOps:
    def __init__(self, total):
        self._total = total
        
        print('creating a new instance')
    
    def __repr__(self):
        return f"ArithmeticOps(total={self._total})"
    
    def __add__(self, value):
        return ArithmeticOps(self._total + value)
    
a = ArithmeticOps(10)    

a + 10

for i in range(1,4):
    print(a+i)
print(a)

 

Notice that we create a new instance each time we add to the instance. This is by convention, and is usually the way Python classes are implemented, now our objects behave the same way as is normal in Python when doing arithmetic. Let's track the value of a for a few iterations of a loop. The output is given below.

 

creating a new instance
ArithmeticOps(total=11)
creating a new instance
ArithmeticOps(total=12)
creating a new instance
ArithmeticOps(total=13)
ArithmeticOps(total=10)

 

What if we don't want to create a new instance each time we add to a class?

 

__iadd__

 

The __iadd__ method represents in-place addition. Its implementation is quite intuitive. 

 

class ArithmeticOps:
    def __init__(self, total):
        self._total = total
        print('creating a new instance')
    
    def __repr__(self):
        return f"ArithmeticOps(total={self._total})"
    
    def __add__(self, value):
        return ArithmeticOps(self._total + value)
        
    def __iadd__(self, value):
        self._total += value
        return self


a = ArithmeticOps(10)   


for i in range(1,4):
    a += i
    print(a)
    
print('after loop a =', a)

 

The output from this script is shown below, notice that we only create a new instance once, all further addition operations simply add to the existing instance. 

 

creating a new instance
ArithmeticOps(total=11)
ArithmeticOps(total=13)
ArithmeticOps(total=16)
after loop a = ArithmeticOps(total=16)

 

Notice from above that we must return self when using in-place addition. We we don't return self then a will = None after the operation is completed, which almost surely is not the desired behavior. 

 

So everything seems to work as expected up to this point. However, let's try the following:

2 + a

Out:
TypeError: unsupported operand type(s) for +: 'int' and 'ArithmeticOps'

 

That may seem quite strange. If we want to implement addition from the right we need to implement another method. 

 

__radd__

 

It should also be mentioned here that using __radd__ is highly dependent on the context in which it is used. Perhaps we just want to return an integer or perhaps we want to return a new object. For the purposes of this tutorial we will just return a new object. 

 

class ArithmeticOps:
    def __init__(self, total):
        self._total = total
        print('creating a new instance')
    
    def __repr__(self):
        return f"ArithmeticOps(total={self._total})"
    
    def __add__(self, value):
        return ArithmeticOps(self._total + value)
        
    def __iadd__(self, value):
        self._total += value
        return self
        
    def __radd__(self, value):
        return ArithmeticOps(value + self._total)
    
    
a = ArithmeticOps(10)         
print(2 + a)

# returns
'''
creating a new instance
creating a new instance
ArithmeticOps(total=12)
'''

 

 

We can use the same protocol for implementing subtraction, multiplication and division. 

 

Subtraction

Standard Subtraction: __sub__

In-place Subtraction: __isub__

Right Subtraction: __rsub__

 

Multiplication

Standard Multiplication: __mul__

In-place Multiplication: __imul__

Right Multiplication: __rmul__

 

Division

Standard Division: __truediv__

In-place Division: __itruediv__

Right Division: __rtruediv__

 

These methods have been implemented below and we will leave it to the reader to test out some examples. 

 

class ArithmeticOps:
    def __init__(self, total):
        self._total = total
        print('creating a new instance')
    
    def __repr__(self):
        return f"ArithmeticOps(total={self._total})"
    
    def __add__(self, value):
        return ArithmeticOps(self._total + value)
        
    def __iadd__(self, value):
        self._total += value
        return self
        
    def __radd__(self, value):
        return ArithmeticOps(value + self._total)
        
    def __sub__(self, value):
        return ArithmeticOps(self._total - value)
              
    def __isub__(self, value):
        self._total -= value
        return self
         
    def __rsub__(self, value):
        return ArithmeticOps(value + self._total)
      
    def __mul__(self, value):
        return ArithmeticOps(self._total * value)
               
    def __imul__(self, value):
        self._total *= value
        return self
    
    def __rmul__(self, value):
        return ArithmeticOps(value * self._total)
               
    def __truediv__(self, value):
        return ArithmeticOps(self._total /  value)
      
    def __itruediv__(self, value):
        self._total /= value
        return self
      
    def __rtruediv__(self, value):
        return value / self._total

a = ArithmeticOps(10)       

 

Real Example

 

(By real here we mean in context.)

It is fair for the reader to be a little confused as to why we would ever need these methods, so let's put the methods into context with a simple example. 

We will create two classes to implement functionality for calculating a restaurant bill. The first class is defined below and named Item, this class takes a name and a price. So for example we would take a steak that costs 29.99 and pass these variables to the constructor.

We will also use two static methods which were discussed here to validate that the name and price a user passes to the constructor are legitimate. 

 

from numbers import Real

class Item:
    def __init__(self, name, price):
        self.name = Item.valid_name(name)
        self.price = Item.valid_price(price)
    
    def __repr__(self):
        return f"Item(name={self.name}, price={self.price})"
    
    @staticmethod
    def valid_name(name):
        if isinstance(name, str) and len(name) > 1:
            return name
        else:
            raise AttributeError('Name must be a string of len > 1')
            
    @staticmethod
    def valid_price(price):
        if isinstance(price, Real) and price > 0:
            return price
        raise AttributeError("price must be a nonnegative real number")
        

s = Item('steak', 29.99)  
        
print(s)
# Item(name=steak, price=29.99)

print(s.name)
# 'steak'

print(s.price)
# 29.99

 

Ok now that we have our Items we will create another class called Bill, this class will keep a total of the items the consumer has had. We need to be able to add items to the bill's total. So consider the methods we have discussed in the examples in the previous section. Clearly __iadd__ is going to need to be implemented. We will also implement the __isub__ method so we can take items away from the Bill. 

We probably want to do some validation in this class also. So we will only accept adding items that are an instance of the Item class we defined above. 

 

class Bill:
    def __init__(self, total=0):
        self.total = 0
        
    def __repr__(self):
        return f"Bill(total={self.total})"
   
    def __iadd__(self, item):
        if isinstance(item, Item):
            print(f"adding {item.name} @ ${item.price} to total")
            self.total += item.price 
            return self
        else:
            raise TypeError('item must be an Item object') 
            
    def __isub__(self, item):
        if isinstance(item, Item):
            print(f"removing {item.name} @ ${item.price} from total")
            self.total += item.price 
            return self
        else:
            raise TypeError('item must be an Item object') 


bill = Bill()            
            
steak = Item('steak',24.99)
pasta = Item('pasta',16.99)
rice = Item('rice',4.0)
bread = Item('bread' , 2.50)          
wine = Item('wine' , 22.50)    
            
bill += steak
bill+= pasta
bill += rice 
bill -= rice
bill += wine

 

Running the script above will result in the following output:

 

adding steak @ $24.99 to total
adding pasta @ $16.99 to total
adding rice @ $4.0 to total
removing rice @ $4.0 from total
adding wine @ $22.5 to total

 

That works as expected. Let's continue to build out our class using the methods we have learned. What would happen if we wanted to add two people's bills together. Say for example the waitress thought two diners were eating separately, but in fact they were together. So for this method we will need to use the __add__ method. In terms of validation we will only allow addition of two bills if the value we try to add is a Bill object. 

 

class Bill:
    def __init__(self, total=0):
        self.total = total
        
    def __repr__(self):
        return f"Bill(total={self.total})"
   
    def __iadd__(self, item):
        if isinstance(item, Item):
            print(f"adding {item.name} @ ${item.price} to total")
            self.total += item.price 
            return self
        else:
            raise TypeError('item must be an Item object') 
            
    def __isub__(self, item):
        if isinstance(item, Item):
            print(f"removing {item.name} @ ${item.price} from total")
            self.total += item.price 
            return self
        else:
            raise TypeError('item must be an Item object') 
            
    def __add__(self, other_bill):
        if isinstance(other_bill, Bill):
            print(f"Adding {other_bill} and {self} = {self.total + other_bill.total}")
            return Bill(self.total + other_bill.total)
        else:
            raise TypeError('Can only add two Bill objects')
   


steak = Item('steak',24.99)
pasta = Item('pasta',16.99)
rice = Item('rice',4.0)
bread = Item('bread' , 2.50)          
wine = Item('wine' , 22.50)   

sally = Bill()
harry = Bill()

sally += steak 
sally += rice
harry += pasta
harry += bread 

new_bill = sally + harry

 

Running this script will result in the following output which we expected.

 

adding steak @ $24.99 to total
adding rice @ $4.0 to total
adding pasta @ $16.99 to total
adding bread @ $2.5 to total
Adding Bill(total=19.49) and Bill(total=28.99) = 48.48

 

Let's add one more method to be able to multiply the total by a nonnegative real number. Putting this into the context of our project, we will make a function below that applies value-added-tax to the final bill. So we need to be able to multiply the final bill by a real number that is greater than 0. Below we also present the entire example from start to finish to avoid any confusion. 

 

from numbers import Real

class Item:
    def __init__(self, name, price):
        self.name = Item.valid_name(name)
        self.price = Item.valid_price(price)
    
    def __repr__(self):
        return f"Item(name={self.name}, price={self.price})"
    
    @staticmethod
    def valid_name(name):
        if isinstance(name, str) and len(name) > 1:
            return name
        else:
            raise AttributeError('Name must be a string of len > 1')
            
    @staticmethod
    def valid_price(price):
        if isinstance(price, Real) and price > 0:
            return price
        raise AttributeError("price must be a nonnegative real number")


class Bill:
    def __init__(self, total=0):
        self.total = total
        
    def __repr__(self):
        return f"Bill(total={self.total})"
   
    def __iadd__(self, item):
        if isinstance(item, Item):
            print(f"adding {item.name} @ ${item.price} to total")
            self.total += item.price 
            return self
        else:
            raise TypeError('item must be an Item object') 
            
    def __isub__(self, item):
        if isinstance(item, Item):
            print(f"removing {item.name} @ ${item.price} from total")
            self.total += item.price 
            return self
        else:
            raise TypeError('item must be an Item object') 
            
    def __add__(self, other_bill):
        if isinstance(other_bill, Bill):
            print(f"Adding {other_bill} and {self} = {self.total + other_bill.total}")
            return Bill(self.total + other_bill.total)
        else:
            raise TypeError('Can only add two Bill objects')
            
    def __mul__(self, value):
        if isinstance(value, Real) and value >=0:
            return Bill(self.total * value)
        else:
            raise TypeError(f"{value} is wrong type, must be nonnegative real number")


steak = Item('steak',24.99)
pasta = Item('pasta',16.99)
rice = Item('rice',4.0)
bread = Item('bread' , 2.50)          
wine = Item('wine' , 22.50)   

sally = Bill()
harry = Bill()

sally += steak 
sally += rice
harry += pasta
harry += bread 

new_bill = sally + harry



def apply_vat(bill_obj, vat_percentage):
    print(f"Bill without tax = {bill_obj}")
    with_vat = bill_obj * (1+ vat_percentage)
    print(f"bill inclusive of value added tax: {with_vat}")
    return with_vat


new_bill = apply_vat(new_bill, 0.15)

 

Clearly there is still a lot of work to do if we wanted to make this a good class. For example, perhaps we wouldn't want to allow subtraction from a bill object if it would make the total less than 0. We will leave this to the reader, Good luck!

 

 

Summary

- We can implement arithmetic operators using the __add__ , __sub__ , __mul__ and __truediv__ methods. These methods usually return a new object whereas the __iadd__ etc results in in-place operations and does not return a new object. 

- Ensure for in-place operations you remember to return self  or Python will return a None object which can be a confusing bug! 

 

 

 

 


Join the discussion

Share this post with your friends!