Improve backward compatibility of Tags APIs

2 files changed, 22 insertions, 9 deletions

diff --git a/python/binaryview.py b/python/binaryview.py
index 1c568789..72dc5213 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -4207,23 +4207,34 @@ class BinaryView:
 		core.BNRemoveTagType(self.handle, tag_type.handle)
 
 	@property
-	def tag_types(self) -> Mapping[str, List['TagType']]:
+	def tag_types(self) -> Mapping[str, Union['TagType', List['TagType']]]:
 		"""
 		``tag_types`` gets a dictionary of all Tag Types present for the view,
 		structured as {Tag Type Name => Tag Type}.
 
+		.. warning:: This method inconsistently returns a list of 'TagType' objects or a single \
+		 'TagType' this behavior will change in future revisions
+
 		:rtype: dict of (str, TagType)
 		"""
 		count = ctypes.c_ulonglong(0)
 		types = core.BNGetTagTypes(self.handle, count)
 		assert types is not None, "core.BNGetTagTypes returned None"
-		result:Mapping[str, List['TagType']] = defaultdict(list)
+		result:Mapping[str, Union['TagType', List['TagType']]] = {}
 		try:
 			for i in range(0, count.value):
 				tag_handle = core.BNNewTagTypeReference(types[i])
 				assert tag_handle is not None, "core.BNNewTagTypeReference returned None"
 				tag = TagType(tag_handle)
-				result[tag.name].append(tag)
+				if tag.name in result:
+					cur_item = result[tag.name]
+					if isinstance(cur_item, list):
+						cur_item.append(tag)
+						result[tag.name] = cur_item
+					else:
+						result[tag.name] = [result[tag.name], tag]
+				else:
+					result[tag.name] = tag
 			return result
 		finally:
 			core.BNFreeTagTypeList(types, count.value)
@@ -4267,7 +4278,7 @@ class BinaryView:
 	def create_auto_tag(self, type:'TagType', data:str) -> 'Tag':
 		return self.create_tag(type, data, False)
 
-	def create_tag(self, type:'TagType', data:str, user:bool=True) -> 'Tag':
+	def create_tag(self, tag_type:'TagType', data:str, user:bool=True) -> 'Tag':
 		"""
 		``create_tag`` creates a new Tag object but does not add it anywhere.
 		Use :py:meth:`create_user_data_tag` to create and add in one step.
@@ -4284,7 +4295,8 @@ class BinaryView:
 			>>> bv.add_user_data_tag(here, tag)
 			>>>
 		"""
-		tag_handle = core.BNCreateTag(type.handle, data)
+		assert isinstance(tag_type, TagType), f"type is not a TagType instead got {type(tag_type)} : {repr(tag_type)}"
+		tag_handle = core.BNCreateTag(tag_type.handle, data)
 		assert tag_handle is not None, "core.BNCreateTag returned None"
 		tag = Tag(tag_handle)
 		core.BNAddTag(self.handle, tag.handle, user)
diff --git a/python/function.py b/python/function.py
index 7d9a6d16..bbab681f 100644
--- a/python/function.py
+++ b/python/function.py
@@ -513,7 +513,7 @@ class Function:
 			arch = self.arch
 		core.BNAddUserAddressTag(self.handle, arch.handle, addr, tag.handle)
 
-	def create_user_address_tag(self, addr:int, type:'binaryview.TagType', data:str, unique:bool=False,
+	def create_user_address_tag(self, addr:int, tag_type:'binaryview.TagType', data:str, unique:bool=False,
 		arch:Optional['architecture.Architecture']=None) -> 'binaryview.Tag':
 		"""
 		``create_user_address_tag`` creates and adds a Tag object at a given
@@ -522,13 +522,14 @@ class Function:
 		inside of a function, use :py:meth:`create_user_data_tag <binaryview.BinaryView.create_user_data_tag>`.
 
 		:param int addr: Address at which to add the tag
-		:param TagType type: Tag Type for the Tag that is created
+		:param TagType tag_type: Tag Type for the Tag that is created
 		:param str data: Additional data for the Tag
 		:param bool unique: If a tag already exists at this location with this data, don't add another
 		:param Architecture arch: Architecture for the block in which the Tag is added (optional)
 		:return: The created Tag
 		:rtype: Tag
 		"""
+		assert isinstance(tag_type, binaryview.TagType), f"type is not a TagType instead got {type(tag_type)} : {repr(tag_type)}"
 		if arch is None:
 			if self.arch is None:
 				raise Exception(f"Can't call {_function_name_()} for function with no architecture specified")
@@ -536,10 +537,10 @@ class Function:
 		if unique:
 			tags = self.get_address_tags_at(addr, arch)
 			for tag in tags:
-				if tag.type == type and tag.data == data:
+				if tag.type == tag_type and tag.data == data:
 					return tag
 
-		tag = self.create_tag(type, data, True)
+		tag = self.create_tag(tag_type, data, True)
 		core.BNAddUserAddressTag(self.handle, arch.handle, addr, tag.handle)
 		return tag