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'