From 949f7c7481395619d2a74cfc6ae7c0a713c89d1c Mon Sep 17 00:00:00 2001
From: Peter LaFosse <peter@vector35.com>
Date: Wed, 29 Sep 2021 16:41:46 -0400
Subject: Fix TypeMapping and SymbolMapping

---
 python/binaryview.py | 45 ++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 3 deletions(-)

(limited to 'python/binaryview.py')

diff --git a/python/binaryview.py b/python/binaryview.py
index 1cf3a62d..f7f253c5 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -1259,6 +1259,7 @@ class Tag:
 class _BinaryViewAssociatedDataStore(associateddatastore._AssociatedDataStore):
 	_defaults = {}
 
+
 class SymbolMapping(collections.abc.Mapping):  # type: ignore
 	def __init__(self, view:'BinaryView'):
 		self._symbol_list = None
@@ -1274,7 +1275,7 @@ class SymbolMapping(collections.abc.Mapping):  # type: ignore
 	def __getitem__(self, key):
 		if self._symbol_cache is None:
 			result = self._view.get_symbols_by_raw_name(key)
-			if result is None:
+			if result == []:
 				raise KeyError(key)
 			return result
 		else:
@@ -1311,6 +1312,37 @@ class SymbolMapping(collections.abc.Mapping):  # type: ignore
 			self._build_symbol_cache()
 		return self._count
 
+	def __contains__(self, value):
+		try:
+			self[value]
+			return True
+		except KeyError:
+			return False
+
+	def keys(self):
+		if self._symbol_cache is None:
+			self._build_symbol_cache()
+		assert self._symbol_cache is not None
+		return self._symbol_cache.keys()
+
+	def items(self):
+		if self._symbol_cache is None:
+			self._build_symbol_cache()
+		assert self._symbol_cache is not None
+		return self._symbol_cache.items()
+
+	def values(self):
+		if self._symbol_cache is None:
+			self._build_symbol_cache()
+		assert self._symbol_cache is not None
+		return self._symbol_cache.values()
+
+	def get(self, value):
+		try:
+			return self[value]
+		except KeyError:
+			return None
+
 
 class TypeMapping(collections.abc.Mapping):  # type: ignore
 	def __init__(self, view:'BinaryView'):
@@ -1358,7 +1390,11 @@ class TypeMapping(collections.abc.Mapping):  # type: ignore
 		return self._count
 
 	def __contains__(self, value):
-		return self[value] != None
+		try:
+			self[value]
+			return True
+		except KeyError:
+			return False
 
 	def __eq__(self, other):
 		return self.view == other.view
@@ -1385,7 +1421,10 @@ class TypeMapping(collections.abc.Mapping):  # type: ignore
 		return self._type_cache.values()
 
 	def get(self, value):
-		return self[value]
+		try:
+			return self[value]
+		except KeyError:
+			return None
 
 
 class FunctionList:
-- 
cgit v1.3.1