Add type hints to basicblock.py, lowlevelil.py, architecture.py

1 files changed, 49 insertions, 29 deletions

diff --git a/python/flowgraph.py b/python/flowgraph.py
index 0407bec0..47ed187c 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -24,20 +24,18 @@ import traceback
 
 # Binary Ninja components
 import binaryninja
-from binaryninja.enums import (BranchType, InstructionTextTokenType, HighlightStandardColor, FlowGraphOption,
+from . import _binaryninjacore as core
+from .enums import (BranchType, InstructionTextTokenType, HighlightStandardColor, FlowGraphOption,
 	EdgePenStyle, ThemeColor)
-from binaryninja import _binaryninjacore as core
-from binaryninja import function
-from binaryninja import binaryview
-from binaryninja import lowlevelil
-from binaryninja import mediumlevelil
-from binaryninja import highlevelil
-from binaryninja import basicblock
-from binaryninja import log
-from binaryninja import highlight
-
-# 2-3 compatibility
-from binaryninja import range
+from . import function
+from . import binaryview
+from . import lowlevelil
+from . import mediumlevelil
+from . import highlevelil
+from . import basicblock
+from . import log
+from . import highlight
+from . import interaction
 
 
 class FlowGraphEdge(object):
@@ -73,8 +71,8 @@ class EdgeStyle(object):
 		result.color = self.color
 		return result
 
-	@classmethod
-	def from_core_struct(cls, edge_style):
+	@staticmethod
+	def from_core_struct(edge_style):
 		return EdgeStyle(edge_style.style, edge_style.width, edge_style.color)
 
 	def __eq__(self, other):
@@ -92,6 +90,7 @@ class FlowGraphNode(object):
 				raise ValueError("flow graph node must be associated with a graph")
 			handle = core.BNCreateFlowGraphNode(graph.handle)
 		self.handle = handle
+		assert self.handle is not None, "core.BNCreateFlowGraphNode returned None"
 		self._graph = graph
 		if self._graph is None:
 			self._graph = FlowGraph(handle = core.BNGetFlowGraphNodeOwner(self.handle))
@@ -113,6 +112,9 @@ class FlowGraphNode(object):
 	def __eq__(self, other):
 		if not isinstance(other, self.__class__):
 			return NotImplemented
+
+		assert self.handle is not None
+		assert other.handle is not None
 		return ctypes.addressof(self.handle.contents) == ctypes.addressof(other.handle.contents)
 
 	def __ne__(self, other):
@@ -121,28 +123,29 @@ class FlowGraphNode(object):
 		return not (self == other)
 
 	def __hash__(self):
+		assert self.handle is not None
 		return hash(ctypes.addressof(self.handle.contents))
 
 	def __iter__(self):
 		count = ctypes.c_ulonglong()
 		lines = core.BNGetFlowGraphNodeLines(self.handle, count)
+		assert lines is not None, "core.BNGetFlowGraphNodeLines returned None"
 		block = self.basic_block
 		try:
 			for i in range(0, count.value):
 				addr = lines[i].addr
 				if (lines[i].instrIndex != 0xffffffffffffffff) and (block is not None) and hasattr(block, 'il_function'):
-					il_instr = block.il_function[lines[i].instrIndex]
+					il_instr = block.__dict__['il_function'][lines[i].instrIndex]
 				else:
 					il_instr = None
-				tokens = function.InstructionTextToken.get_instruction_lines(lines[i].tokens, lines[i].count)
+				tokens = function.InstructionTextToken._from_core_struct(lines[i].tokens, lines[i].count)
 				yield function.DisassemblyTextLine(tokens, addr, il_instr)
 		finally:
 			core.BNFreeDisassemblyTextLines(lines, count.value)
 
 	@property
 	def graph(self):
-		""" """
-		return self._graph
+				return self._graph
 
 	@graph.setter
 	def graph(self, value):
@@ -166,8 +169,8 @@ class FlowGraphNode(object):
 			block = lowlevelil.LowLevelILBasicBlock(view, block,
 				lowlevelil.LowLevelILFunction(func.arch, core.BNGetBasicBlockLowLevelILFunction(block), func))
 		elif core.BNIsMediumLevelILBasicBlock(block):
-			block = mediumlevelil.MediumLevelILBasicBlock(view, block,
-				mediumlevelil.MediumLevelILFunction(func.arch, core.BNGetBasicBlockMediumLevelILFunction(block), func))
+			mlil_func = mediumlevelil.MediumLevelILFunction(func.arch, core.BNGetBasicBlockMediumLevelILFunction(block), func)
+			block = mediumlevelil.MediumLevelILBasicBlock(block, mlil_func, view)
 		else:
 			block = basicblock.BasicBlock(block, view)
 		return block
@@ -204,16 +207,17 @@ class FlowGraphNode(object):
 		"""Flow graph block list of text lines"""
 		count = ctypes.c_ulonglong()
 		lines = core.BNGetFlowGraphNodeLines(self.handle, count)
+		assert lines is not None, "core.BNGetFlowGraphNodeLines returned None"
 		block = self.basic_block
 		result = []
 		for i in range(0, count.value):
 			addr = lines[i].addr
 			if (lines[i].instrIndex != 0xffffffffffffffff) and (block is not None) and hasattr(block, 'il_function'):
-				il_instr = block.il_function[lines[i].instrIndex]
+				il_instr = block.__dict__['il_function'][lines[i].instrIndex]
 			else:
 				il_instr = None
 			color = highlight.HighlightColor._from_core_struct(lines[i].highlight)
-			tokens = function.InstructionTextToken.get_instruction_lines(lines[i].tokens, lines[i].count)
+			tokens = function.InstructionTextToken._from_core_struct(lines[i].tokens, lines[i].count)
 			result.append(function.DisassemblyTextLine(tokens, addr, il_instr, color))
 		core.BNFreeDisassemblyTextLines(lines, count.value)
 		return result
@@ -247,7 +251,7 @@ class FlowGraphNode(object):
 				color = highlight.HighlightColor(color)
 			line_buf[i].highlight = color._get_core_struct()
 			line_buf[i].count = len(line.tokens)
-			line_buf[i].tokens = function.InstructionTextToken.get_instruction_lines(line.tokens)
+			line_buf[i].tokens = function.InstructionTextToken._get_core_struct(line.tokens)
 		core.BNSetFlowGraphNodeLines(self.handle, line_buf, len(lines))
 
 	@property
@@ -255,6 +259,7 @@ class FlowGraphNode(object):
 		"""Flow graph block list of outgoing edges (read-only)"""
 		count = ctypes.c_ulonglong()
 		edges = core.BNGetFlowGraphNodeOutgoingEdges(self.handle, count)
+		assert edges is not None, "core.BNGetFlowGraphNodeOutgoingEdges returned None"
 		result = []
 		for i in range(0, count.value):
 			branch_type = BranchType(edges[i].type)
@@ -273,6 +278,7 @@ class FlowGraphNode(object):
 		"""Flow graph block list of incoming edges (read-only)"""
 		count = ctypes.c_ulonglong()
 		edges = core.BNGetFlowGraphNodeIncomingEdges(self.handle, count)
+		assert edges is not None, "core.BNGetFlowGraphNodeIncomingEdges returned None"
 		result = []
 		for i in range(0, count.value):
 			branch_type = BranchType(edges[i].type)
@@ -419,6 +425,8 @@ class FlowGraph(object):
 	def __eq__(self, other):
 		if not isinstance(other, self.__class__):
 			return NotImplemented
+		assert self.handle is not None
+		assert other.handle is not None
 		return ctypes.addressof(self.handle.contents) == ctypes.addressof(other.handle.contents)
 
 	def __ne__(self, other):
@@ -427,6 +435,7 @@ class FlowGraph(object):
 		return not (self == other)
 
 	def __hash__(self):
+		assert self.handle is not None
 		return hash(ctypes.addressof(self.handle.contents))
 
 	def __setattr__(self, name, value):
@@ -438,9 +447,12 @@ class FlowGraph(object):
 	def __iter__(self):
 		count = ctypes.c_ulonglong()
 		nodes = core.BNGetFlowGraphNodes(self.handle, count)
+		assert nodes is not None, "core.BNGetFlowGraphNodes returned None"
 		try:
 			for i in range(0, count.value):
-				yield FlowGraphNode(self, core.BNNewFlowGraphNodeReference(nodes[i]))
+				fg_node = core.BNNewFlowGraphNodeReference(nodes[i])
+				assert fg_node is not None, "core.BNNewFlowGraphNodeReference returned None"
+				yield FlowGraphNode(self, fg_node)
 		finally:
 			core.BNFreeFlowGraphNodeList(nodes, count.value)
 
@@ -473,7 +485,9 @@ class FlowGraph(object):
 			graph = self.update()
 			if graph is NotImplemented:
 				return None
-			return ctypes.cast(core.BNNewFlowGraphReference(graph.handle), ctypes.c_void_p).value
+			flow_graph = core.BNNewFlowGraphReference(graph.handle)
+			assert flow_graph is not None, "core.BNNewFlowGraphReference returned None"
+			return ctypes.cast(flow_graph, ctypes.c_void_p).value
 		except:
 			log.log_error(traceback.format_exc())
 			return None
@@ -557,9 +571,12 @@ class FlowGraph(object):
 		"""List of nodes in graph (read-only)"""
 		count = ctypes.c_ulonglong()
 		blocks = core.BNGetFlowGraphNodes(self.handle, count)
+		assert blocks is not None, "core.BNGetFlowGraphNodes returned None"
 		result = []
 		for i in range(0, count.value):
-			result.append(FlowGraphNode(self, core.BNNewFlowGraphNodeReference(blocks[i])))
+			fg_node = core.BNNewFlowGraphNodeReference(blocks[i])
+			assert fg_node is not None, "core.BNNewFlowGraphNodeReference returned None"
+			result.append(FlowGraphNode(self, fg_node))
 		core.BNFreeFlowGraphNodeList(blocks, count.value)
 		return result
 
@@ -747,9 +764,12 @@ class FlowGraph(object):
 	def get_nodes_in_region(self, left, top, right, bottom):
 		count = ctypes.c_ulonglong()
 		nodes = core.BNGetFlowGraphNodesInRegion(self.handle, left, top, right, bottom, count)
+		assert nodes is not None, "core.BNGetFlowGraphNodesInRegion returned None"
 		result = []
 		for i in range(0, count.value):
-			result.append(FlowGraphNode(self, core.BNNewFlowGraphNodeReference(nodes[i])))
+			fg_node = core.BNNewFlowGraphNodeReference(nodes[i])
+			assert fg_node is not None, "core.BNNewFlowGraphNodeReference returned None"
+			result.append(FlowGraphNode(self, fg_node))
 		core.BNFreeFlowGraphNodeList(nodes, count.value)
 		return result
 
@@ -769,7 +789,7 @@ class FlowGraph(object):
 
 		:param str title: Title to show in the new tab
 		"""
-		binaryninja.interaction.show_graph_report(title, self)
+		interaction.show_graph_report(title, self)
 
 	def update(self):
 		"""