CODEWITHSUNDEEP
pythonpython-3.x
user10019227
user10019227

Reputation: 117

Using functools.total_ordering for comparisons

I have the below code...

class BankAccount:
    """ Simple BankAccount class """

    def __init__(self, balance=0):
        """Initialize account with balance"""
        self.balance = balance

    def deposit(self, amount):
        """Deposit amount to this account"""
        self.balance += amount

    def withdraw(self, amount):
        """Withdraw amount from this account"""
        self.balance -= amount

    def __str__(self):
        return 'Account with a balance of {}'.format(self.balance)

    def __repr__(self):
        return "BankAccount(balance={})".format(self.balance)

    def __bool__(self):
        if self.balance > 0:
            return True
        else:
            return False

The code is basically a simple bank account simulator. I want to implement comparisons for BankAccount objects, such that instances can be compared based on their balance. I want to do this using functools.total_ordering. An expected output is below...

    account1 = BankAccount(100)
    account2 = BankAccount()
    account3 = BankAccount(100)
    account1 == account2
False
    account1 == account3
True
    account1 != account3
False
    account1 < account2
False
    account1 >= account2
True

How would I do this?

Upvotes: 1

Views: 2686

Answers (1)

CodeSamurai777
CodeSamurai777

Reputation: 3355

You just need to define at least one of the functions:

__lt__(), __le__(), __gt__(), or __ge__() In addition, the class should supply an __eq__() method.

Then you use the decorator like so:

from functools import total_ordering

@total_ordering
class BankAccount:
""" Simple BankAccount class """

   def __init__(self, balance=0):
    """Initialize account with balance"""
       self.balance = balance
   def __lt__(self, other):
       return self.balance  < other.balance 
   def __eq__(self,other):
       return self.balance == other.balance

Upvotes: 7

Related Questions