Home > Blockchain >  Is there a way to make comparsion operation between two objects with decorators in python?
Is there a way to make comparsion operation between two objects with decorators in python?

Time:10-04

What I want is to make comparsion between two objects with decorator. Something like this

@some_decorator
class SomeClass:
    def __init__(self, value):
        self.value = value


s1 = SomeClass(10)
s2 = SomeClass(5)
print(s1 > s2)  # True

I know about magic methods. I am just curios about is it really possible or not

CodePudding user response:

You can define a couple comparisons function in the decorator function and then add them to the class. To take an argument, you can wrap the function. Using functools.total_ordering is a convenient way to add other comparison functions without needing to write them all out (note the caveat in the docs about performance if that's important).

from functools import total_ordering

def compare(attribute):
    def deco(klass):
        def __eq__(self, other):
            if isinstance(other, klass):
                return getattr(self, attribute) == getattr(other, attribute)
            else:
                return NotImplemented

        def __gt__(self, other):
            if isinstance(other, klass):
                return getattr(self, attribute) > getattr(other, attribute)
            else:
                return NotImplemented

        klass.__eq__ = __eq__ 
        klass.__gt__ = __gt__

        return total_ordering(klass)
    
    return deco

@compare('value')
class SomeClass():
    def __init__(self, value):
        self.value = value


s1 = SomeClass(10)
s2 = SomeClass(5)
print(s1 > s2)  # True
print(s1 >= s2)  # True
print(s1 < s2)  # False
print(s1 <= s2)  # False
print(s1 == s2) # False
print(s1 != s2) # True

As @juanpa.arrivillaga recommends in the comments, this checks the instance so a comparison like the following behaves as expected:

s1 > 50
# TypeError: '>' not supported between instances of 'SomeClass' and 'int'

Instead of the more confusing: AttributeError: 'int' object has no attribute 'value'

  • Related