Various Python 3 support changes

1 files changed, 18 insertions, 7 deletions

diff --git a/python/function.py b/python/function.py
index 6392d7f2..8407eece 100644
--- a/python/function.py
+++ b/python/function.py
@@ -23,20 +23,20 @@ import threading
 import traceback
 import ctypes
 
-# Binary Ninja components -- additional imports belong in the appropriate class
+# Binary Ninja components
 import binaryninja
 from binaryninja import _binaryninjacore as core
+from binaryninja import associateddatastore  #required in the main scope due to being an argument for _FunctionAssociatedDataStore
+from binaryninja import highlight
+from binaryninja import log
+from binaryninja import types
 from binaryninja.enums import (AnalysisSkipReason, FunctionGraphType, BranchType, SymbolType, InstructionTextTokenType,
 	HighlightStandardColor, HighlightColorStyle, RegisterValueType, ImplicitRegisterExtend,
 	DisassemblyOption, IntegerDisplayType, InstructionTextTokenContext, VariableSourceType,
 	FunctionAnalysisSkipOverride)
-from binaryninja import associateddatastore #required in the main scope due to being an argument for _FunctionAssociatedDataStore
-from binaryninja import types
-from binaryninja import highlight
-from binaryninja import log
 
 # 2-3 compatibility
-from six.moves import range
+from binaryninja import range
 
 
 class LookupTableEntry(object):
@@ -49,7 +49,7 @@ class LookupTableEntry(object):
 
 
 class RegisterValue(object):
-	def __init__(self, arch = None, value = None, confidence = binaryninja.types.max_confidence):
+	def __init__(self, arch = None, value = None, confidence = types.max_confidence):
 		self.is_constant = False
 		if value is None:
 			self.type = RegisterValueType.UndeterminedValue
@@ -1160,6 +1160,8 @@ class Function(object):
 			for j in range(0, lines[i].count):
 				token_type = InstructionTextTokenType(lines[i].tokens[j].type)
 				text = lines[i].tokens[j].text
+				if not isinstance(text, str):
+					text = text.encode("charmap")
 				value = lines[i].tokens[j].value
 				size = lines[i].tokens[j].size
 				operand = lines[i].tokens[j].operand
@@ -1836,6 +1838,7 @@ class DisassemblySettings(object):
 		core.BNSetDisassemblySettingsOption(self.handle, option, state)
 
 
+_pending_function_graph_completion_events = {}
 class FunctionGraph(object):
 	def __init__(self, view, handle):
 		self.view = view
@@ -1969,6 +1972,9 @@ class FunctionGraph(object):
 		try:
 			if self._on_complete is not None:
 				self._on_complete()
+				global _pending_function_graph_completion_events
+				if id(self) in _pending_function_graph_completion_events:
+					del _pending_function_graph_completion_events[id(self)]
 		except:
 			log.log_error(traceback.format_exc())
 
@@ -1993,11 +1999,16 @@ class FunctionGraph(object):
 		self._wait_cond.release()
 
 	def on_complete(self, callback):
+		global _pending_function_graph_completion_events
+		_pending_function_graph_completion_events[id(self)] = self
 		self._on_complete = callback
 		core.BNSetFunctionGraphCompleteCallback(self.handle, None, self._cb)
 
 	def abort(self):
 		core.BNAbortFunctionGraph(self.handle)
+		global _pending_function_graph_completion_events
+		if id(self) in _pending_function_graph_completion_events:
+			del _pending_function_graph_completion_events[id(self)]
 
 	def get_blocks_in_region(self, left, top, right, bottom):
 		count = ctypes.c_ulonglong()