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

1 files changed, 34 insertions, 28 deletions

diff --git a/python/callingconvention.py b/python/callingconvention.py
index c29404aa..096a6b73 100644
--- a/python/callingconvention.py
+++ b/python/callingconvention.py
@@ -22,13 +22,12 @@ import traceback
 import ctypes
 
 # Binary Ninja components
-import binaryninja
-from binaryninja import _binaryninjacore as core
-from binaryninja import log
-from binaryninja.enums import VariableSourceType
-
-# 2-3 compatibility
-from binaryninja import range
+from . import _binaryninjacore as core
+from .enums import VariableSourceType
+from . import log
+from . import variable
+from . import function
+from . import architecture
 
 
 class CallingConvention(object):
@@ -50,7 +49,7 @@ class CallingConvention(object):
 
 	_registered_calling_conventions = []
 
-	def __init__(self, arch=None, name=None, handle=None, confidence=binaryninja.types.max_confidence):
+	def __init__(self, arch=None, name=None, handle=None, confidence=core.max_confidence):
 		if handle is None:
 			if arch is None or name is None:
 				self.handle = None
@@ -82,7 +81,7 @@ class CallingConvention(object):
 			self.__class__._registered_calling_conventions.append(self)
 		else:
 			self.handle = handle
-			self.arch = binaryninja.architecture.CoreArchitecture._from_cache(core.BNGetCallingConventionArchitecture(self.handle))
+			self.arch = architecture.CoreArchitecture._from_cache(core.BNGetCallingConventionArchitecture(self.handle))
 			self.__dict__["name"] = core.BNGetCallingConventionName(self.handle)
 			self.__dict__["arg_regs_share_index"] = core.BNAreArgumentRegistersSharedIndex(self.handle)
 			self.__dict__["arg_regs_for_varargs"] = core.BNAreArgumentRegistersUsedForVarArgs(self.handle)
@@ -92,6 +91,7 @@ class CallingConvention(object):
 
 			count = ctypes.c_ulonglong()
 			regs = core.BNGetCallerSavedRegisters(self.handle, count)
+			assert regs is not None, "core.BNGetCallerSavedRegisters returned None"
 			result = []
 			arch = self.arch
 			for i in range(0, count.value):
@@ -101,6 +101,7 @@ class CallingConvention(object):
 
 			count = ctypes.c_ulonglong()
 			regs = core.BNGetCalleeSavedRegisters(self.handle, count)
+			assert regs is not None, "core.BNGetCalleeSavedRegisters returned None"
 			result = []
 			arch = self.arch
 			for i in range(0, count.value):
@@ -110,6 +111,7 @@ class CallingConvention(object):
 
 			count = ctypes.c_ulonglong()
 			regs = core.BNGetIntegerArgumentRegisters(self.handle, count)
+			assert regs is not None, "core.BNGetIntegerArgumentRegisters returned None"
 			result = []
 			arch = self.arch
 			for i in range(0, count.value):
@@ -119,6 +121,7 @@ class CallingConvention(object):
 
 			count = ctypes.c_ulonglong()
 			regs = core.BNGetFloatArgumentRegisters(self.handle, count)
+			assert regs is not None, "core.BNGetFloatArgumentRegisters returned None"
 			result = []
 			arch = self.arch
 			for i in range(0, count.value):
@@ -152,6 +155,7 @@ class CallingConvention(object):
 
 			count = ctypes.c_ulonglong()
 			regs = core.BNGetImplicitlyDefinedRegisters(self.handle, count)
+			assert regs is not None, "core.BNGetImplicitlyDefinedRegisters returned None"
 			result = []
 			arch = self.arch
 			for i in range(0, count.value):
@@ -174,6 +178,8 @@ class CallingConvention(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):
@@ -182,6 +188,7 @@ class CallingConvention(object):
 		return not (self == other)
 
 	def __hash__(self):
+		assert self.handle is not None
 		return hash(ctypes.addressof(self.handle.contents))
 
 	def _get_caller_saved_regs(self, ctxt, count):
@@ -308,7 +315,7 @@ class CallingConvention(object):
 		try:
 			if self.__class__.float_return_reg is None:
 				return 0xffffffff
-			return self.arch.regs[self.__class__.float_int_return_reg].index
+			return self.arch.regs[self.__class__.float_return_reg].index
 		except:
 			log.log_error(traceback.format_exc())
 			return False
@@ -339,23 +346,23 @@ class CallingConvention(object):
 
 	def _get_incoming_reg_value(self, ctxt, reg, func, result):
 		try:
-			func_obj = binaryninja.function.Function(handle = core.BNNewFunctionReference(func))
+			func_obj = function.Function(handle = core.BNNewFunctionReference(func))
 			reg_name = self.arch.get_reg_name(reg)
 			api_obj = self.perform_get_incoming_reg_value(reg_name, func_obj)._to_api_object()
 		except:
 			log.log_error(traceback.format_exc())
-			api_obj = binaryninja.function.RegisterValue()._to_api_object()
+			api_obj = variable.RegisterValue()._to_api_object()
 		result[0].state = api_obj.state
 		result[0].value = api_obj.value
 
 	def _get_incoming_flag_value(self, ctxt, reg, func, result):
 		try:
-			func_obj = binaryninja.function.Function(handle = core.BNNewFunctionReference(func))
+			func_obj = function.Function(handle = core.BNNewFunctionReference(func))
 			reg_name = self.arch.get_reg_name(reg)
 			api_obj = self.perform_get_incoming_flag_value(reg_name, func_obj)._to_api_object()
 		except:
 			log.log_error(traceback.format_exc())
-			api_obj = binaryninja.function.RegisterValue()._to_api_object()
+			api_obj = variable.RegisterValue()._to_api_object()
 		result[0].state = api_obj.state
 		result[0].value = api_obj.value
 
@@ -364,8 +371,8 @@ class CallingConvention(object):
 			if func is None:
 				func_obj = None
 			else:
-				func_obj = binaryninja.function.Function(handle = core.BNNewFunctionReference(func))
-			in_var_obj = binaryninja.function.Variable(func_obj, in_var[0].type, in_var[0].index, in_var[0].storage)
+				func_obj = function.Function(handle = core.BNNewFunctionReference(func))
+			in_var_obj = variable.Variable(func_obj, in_var[0].type, in_var[0].index, in_var[0].storage)
 			out_var = self.perform_get_incoming_var_for_parameter_var(in_var_obj, func_obj)
 			result[0].type = out_var.source_type
 			result[0].index = out_var.index
@@ -381,8 +388,8 @@ class CallingConvention(object):
 			if func is None:
 				func_obj = None
 			else:
-				func_obj = binaryninja.function.Function(handle = core.BNNewFunctionReference(func))
-			in_var_obj = binaryninja.function.Variable(func_obj, in_var[0].type, in_var[0].index, in_var[0].storage)
+				func_obj = function.Function(handle = core.BNNewFunctionReference(func))
+			in_var_obj = variable.Variable(func_obj, in_var[0].type, in_var[0].index, in_var[0].storage)
 			out_var = self.perform_get_parameter_var_for_incoming_var(in_var_obj, func_obj)
 			result[0].type = out_var.source_type
 			result[0].index = out_var.index
@@ -397,11 +404,11 @@ class CallingConvention(object):
 		reg_stack = self.arch.get_reg_stack_for_reg(reg)
 		if reg_stack is not None:
 			if reg == self.arch.reg_stacks[reg_stack].stack_top_reg:
-				return binaryninja.function.RegisterValue.constant(0)
-		return binaryninja.function.RegisterValue()
+				return variable.RegisterValue.constant(0)
+		return variable.RegisterValue()
 
 	def perform_get_incoming_flag_value(self, reg, func):
-		return binaryninja.function.RegisterValue()
+		return variable.RegisterValue()
 
 	def perform_get_incoming_var_for_parameter_var(self, in_var, func):
 		in_buf = core.BNVariable()
@@ -412,7 +419,7 @@ class CallingConvention(object):
 		name = None
 		if (func is not None) and (out_var.type == VariableSourceType.RegisterVariableSourceType):
 			name = func.arch.get_reg_name(out_var.storage)
-		return binaryninja.function.Variable(func, out_var.type, out_var.index, out_var.storage, name)
+		return variable.Variable(func, out_var.type, out_var.index, out_var.storage, name)
 
 	def perform_get_parameter_var_for_incoming_var(self, in_var, func):
 		in_buf = core.BNVariable()
@@ -420,7 +427,7 @@ class CallingConvention(object):
 		in_buf.index = in_var.index
 		in_buf.storage = in_var.storage
 		out_var = core.BNGetDefaultParameterVariableForIncomingVariable(self.handle, in_buf)
-		return binaryninja.function.Variable(func, out_var.type, out_var.index, out_var.storage)
+		return variable.Variable(func, out_var.type, out_var.index, out_var.storage)
 
 	def with_confidence(self, confidence):
 		return CallingConvention(self.arch, handle = core.BNNewCallingConventionReference(self.handle),
@@ -431,14 +438,14 @@ class CallingConvention(object):
 		func_handle = None
 		if func is not None:
 			func_handle = func.handle
-		return binaryninja.function.RegisterValue(self.arch, core.BNGetIncomingRegisterValue(self.handle, reg_num, func_handle))
+		return variable.RegisterValue(self.arch, core.BNGetIncomingRegisterValue(self.handle, reg_num, func_handle))
 
 	def get_incoming_flag_value(self, flag, func):
 		reg_num = self.arch.get_flag_index(flag)
 		func_handle = None
 		if func is not None:
 			func_handle = func.handle
-		return binaryninja.function.RegisterValue(self.arch, core.BNGetIncomingFlagValue(self.handle, reg_num, func_handle))
+		return variable.RegisterValue(self.arch, core.BNGetIncomingFlagValue(self.handle, reg_num, func_handle))
 
 	def get_incoming_var_for_parameter_var(self, in_var, func):
 		in_buf = core.BNVariable()
@@ -453,7 +460,7 @@ class CallingConvention(object):
 		name = None
 		if (func is not None) and (out_var.type == VariableSourceType.RegisterVariableSourceType):
 			name = func.arch.get_reg_name(out_var.storage)
-		return binaryninja.function.Variable(func, out_var.type, out_var.index, out_var.storage, name)
+		return variable.Variable(func, out_var.type, out_var.index, out_var.storage, name)
 
 	def get_parameter_var_for_incoming_var(self, in_var, func):
 		in_buf = core.BNVariable()
@@ -465,11 +472,10 @@ class CallingConvention(object):
 		else:
 			func_obj = func.handle
 		out_var = core.BNGetParameterVariableForIncomingVariable(self.handle, in_buf, func_obj)
-		return binaryninja.function.Variable(func, out_var.type, out_var.index, out_var.storage)
+		return variable.Variable(func, out_var.type, out_var.index, out_var.storage)
 
 	@property
 	def arch(self):
-		""" """
 		return self._arch
 
 	@arch.setter