From 8f1bc5944b849f3ad0cbbd5972d5ab9948ed4d57 Mon Sep 17 00:00:00 2001
From: Peter LaFosse <peter@vector35.com>
Date: Sat, 2 May 2020 11:01:35 -0400
Subject: Fix up equality operators and had hash operators

---
 python/architecture.py | 95 ++++++++++++++++++++++++++++++--------------------
 1 file changed, 57 insertions(+), 38 deletions(-)

(limited to 'python/architecture.py')

diff --git a/python/architecture.py b/python/architecture.py
index ad295c91..8495255c 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -79,12 +79,6 @@ class _ArchitectureMetaClass(type):
 		cls._registered_cb = arch._cb
 		arch.handle = core.BNRegisterArchitecture(cls.name, arch._cb)
 
-	def __setattr__(self, name, value):
-		try:
-			type.__setattr__(self, name, value)
-		except AttributeError:
-			raise AttributeError("attribute '%s' is read only" % name)
-
 
 class Architecture(with_metaclass(_ArchitectureMetaClass, object)):
 	"""
@@ -387,15 +381,32 @@ class Architecture(with_metaclass(_ArchitectureMetaClass, object)):
 		self._pending_name_and_type_lists = {}
 		self._pending_type_lists = {}
 
-	def __eq__(self, value):
-		if not isinstance(value, Architecture):
-			return False
-		return ctypes.addressof(self.handle.contents) == ctypes.addressof(value.handle.contents)
+	def __repr__(self):
+		return "<arch: %s>" % self.name
 
-	def __ne__(self, value):
-		if not isinstance(value, Architecture):
-			return True
-		return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents)
+	def __eq__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return ctypes.addressof(self.handle.contents) == ctypes.addressof(other.handle.contents)
+
+	def __ne__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return not (self == other)
+
+	def __hash__(self):
+		return hash(ctypes.addressof(self.handle.contents))
+
+	def __setattr__(self, name, value):
+		if ((name == "name") or (name == "endianness") or (name == "address_size") or
+			(name == "default_int_size") or (name == "regs") or (name == "get_max_instruction_length") or
+			(name == "get_instruction_alignment")):
+			raise AttributeError("attribute '%s' is read only" % name)
+		else:
+			try:
+				object.__setattr__(self, name, value)
+			except AttributeError:
+				raise AttributeError("attribute '%s' is read only" % name)
 
 	@property
 	def list(self):
@@ -442,21 +453,6 @@ class Architecture(with_metaclass(_ArchitectureMetaClass, object)):
 		core.BNFreeTypeLibraryList(handles, count.value)
 		return result
 
-
-	def __setattr__(self, name, value):
-		if ((name == "name") or (name == "endianness") or (name == "address_size") or
-			(name == "default_int_size") or (name == "regs") or (name == "get_max_instruction_length") or
-			(name == "get_instruction_alignment")):
-			raise AttributeError("attribute '%s' is read only" % name)
-		else:
-			try:
-				object.__setattr__(self, name, value)
-			except AttributeError:
-				raise AttributeError("attribute '%s' is read only" % name)
-
-	def __repr__(self):
-		return "<arch: %s>" % self.name
-
 	def _init(self, ctxt, handle):
 		self.handle = handle
 
@@ -2749,15 +2745,38 @@ 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
+	def __eq__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return (self.function, self.arch, self.address) == (other.address, other.function, other.arch)
+
+	def __ne__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return not (self == other)
+
+	def __lt__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return self.address < other.address
+
+	def __gt__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return self.address > other.address
+
+	def __gt__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return self.address >= other.address
+
+	def __le__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return self.address <= other.address
+
+	def __hash__(self):
+		hash((self._function, self._arch, self._address))
 
 	@property
 	def function(self):
-- 
cgit v1.3.1