add __lt__ to referencesource objects and raise exception if invalid type

2 files changed, 11 insertions, 1 deletions

diff --git a/python/architecture.py b/python/architecture.py
index b9001bf3..a6e1a083 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -2724,6 +2724,16 @@ class ReferenceSource(object):
 		else:
 			return "<ref: %#x>" % self._address
 
+	def __eq__(self, value):
+		if not isinstance(value, ReferenceSource):
+			return False
+		return self.function == value.function and self.arch == value.arch and self.address == value.address
+
+	def __lt__(self, value):
+		if not isinstance(value, ReferenceSource):
+			raise TypeError("Can only compare to other ReferenceSource objects")
+		return self.address < value.address
+
 	@property
 	def function(self):
 		""" """
diff --git a/python/function.py b/python/function.py
index e0d179b2..bdd04590 100644
--- a/python/function.py
+++ b/python/function.py
@@ -752,7 +752,7 @@ class Function(object):
 
 	def __lt__(self, value):
 		if not isinstance(value, Function):
-			return False
+			raise TypeError("Can only compare to other Function objects")
 		return self.start < value.start
 
 	def __eq__(self, value):