From b67f98d1e982b37ab80c32c019c5c221512e608e Mon Sep 17 00:00:00 2001 From: Jordan Wiens Date: Mon, 16 Mar 2020 21:38:29 -0400 Subject: add __lt__ to referencesource objects and raise exception if invalid type --- python/architecture.py | 10 ++++++++++ python/function.py | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) (limited to 'python') 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 "" % 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): -- cgit v1.3.1