Fix for issue with custom display types

1 files changed, 21 insertions, 5 deletions

diff --git a/python/datarender.py b/python/datarender.py
index 4b81b1dd..52591ebc 100644
--- a/python/datarender.py
+++ b/python/datarender.py
@@ -32,6 +32,22 @@ from binaryninja import log
 from binaryninja import types
 from binaryninja import highlight
 
+
+class TypeContext(object):
+	def __init__(self, _type, _offset):
+		self._type = _type
+		self._offset = _offset
+
+	@property
+	def type(self):
+		"""The Type object for the current context record"""
+		return self._type
+
+	@property
+	def offset(self):
+		"""The offset into the given type object"""
+		return self._offset
+
 class DataRenderer(object):
 	"""
 	DataRenderer objects tell the Linear View how to render specific types.
@@ -61,7 +77,7 @@ class DataRenderer(object):
 				return [DisassemblyTextLine(prefix, addr)]
 			def __del__(self):
 				pass
-		
+
 		BarDataRenderer().register_type_specific()
 
 	Note that the formatting is sub-optimal to work around an issue with Sphinx and reStructured text
@@ -79,8 +95,8 @@ class DataRenderer(object):
 	@classmethod
 	def is_type_of_struct_name(cls, type, name, context):
 		return (type.type_class == enums.TypeClass.StructureTypeClass and len(context) > 0
-			and context[-1].type_class == enums.TypeClass.NamedTypeReferenceClass and
-			context[-1].named_type_reference.name == name)
+			and context[-1].type.type_class == enums.TypeClass.NamedTypeReferenceClass and
+			context[-1].type.named_type_reference.name == name)
 
 	def register_type_specific(self):
 		core.BNRegisterTypeSpecificDataRenderer(core.BNGetDataRendererContainer(), self.handle)
@@ -103,7 +119,7 @@ class DataRenderer(object):
 			type = types.Type(handle=core.BNNewTypeReference(type))
 			pycontext = []
 			for i in range(0, ctxCount):
-				pycontext.append(types.Type(core.BNNewTypeReference(context[i])))
+				pycontext.append(TypeContext(types.Type(core.BNNewTypeReference(context[i].type)), context[i].offset))
 			return self.perform_is_valid_for_data(ctxt, view, addr, type, pycontext)
 		except:
 			log.log_error(traceback.format_exc())
@@ -118,7 +134,7 @@ class DataRenderer(object):
 			prefixTokens = function.InstructionTextToken.get_instruction_lines(prefix, prefixCount)
 			pycontext = []
 			for i in range(ctxCount):
-				pycontext.append(types.Type(core.BNNewTypeReference(typeCtx[i])))
+				pycontext.append(TypeContext(types.Type(core.BNNewTypeReference(typeCtx[i].type)), typeCtx[i].offset))
 
 			result = self.perform_get_lines_for_data(ctxt, view, addr, type, prefixTokens, width, pycontext)