Friday, March 6, 2009

Comparing and the __cmp__ method in Python

Sometimes it would be nice to be able to ask questions such as “does this array hold an object that has its instance variable ‘foo’ set to 42?” in a simpler way than by iterating through the entire array whilst comparing each objects ‘foo’ attribute to 42. This can be done quite easily in Python by overloading the __cmp__ method in your classes. The following example shows how that works in practice:

class Foo(object):
def __init__(self, a_string, an_integer):
self.my_string = a_string
self.my_integer = an_integer

def __cmp__(self, other):
if type(other) == type(self):
return cmp(self.my_string, other.my_string) and\
cmp(self.my_integer, other.my_integer)
elif type(other) == int:
return cmp(self.my_integer, other)
elif type(other) == str:
return cmp(self.my_string, other)
else:
raise TypeError('Unable to compare %s to %s'%(type(self), type(other)))

f1 = Foo('hello', 1)
f2 = Foo('world', 2)
f3 = Foo('!', 3)
my_list = [f1, f2]
print f1 in my_list # Prints "True"
print f3 in my_list # Prints "False"
print 1 in my_list # Prints "True"
print 42 in my_list # Prints "False"
print 'hello' in my_list # Prints "True"
print 'mojn' in my_list # Prints "False"
try:
(1+0j) in my_list # Raises a TypeError
except TypeError, e:
print e.message

No comments:

Post a Comment