4 files changed, 44 insertions, 10 deletions
diff --git a/python/function.py b/python/function.py
index 41b5ac82..152450e8 100644
--- a/python/function.py
+++ b/python/function.py
@@ -170,7 +170,7 @@ class VariableReferenceSource:
 
 
 class BasicBlockList:
-	def __init__(self, function:'Function'):
+	def __init__(self, function:Union['Function', 'lowlevelil.LowLevelILFunction', 'mediumlevelil.MediumLevelILFunction', 'highlevelil.HighLevelILFunction']):
 		self._count, self._blocks = function._basic_block_list()
 		self._function = function
 		self._n = 0
@@ -218,6 +218,43 @@ class BasicBlockList:
 		raise ValueError("BasicBlockList.__getitem__ supports argument of type integer or slice only")
 
 
+
+class LowLevelILBasicBlockList(BasicBlockList):
+
+	def __repr__(self):
+		return f"<LowLevelILBasicBlockList {len(self)} BasicBlocks: {list(self)}>"
+
+	def __getitem__(self, i:Union[int, slice]) -> Union['lowlevelil.LowLevelILBasicBlock', List['lowlevelil.LowLevelILBasicBlock']]:
+		return BasicBlockList.__getitem__(self, i) # type: ignore
+
+	def __next__(self) -> 'lowlevelil.LowLevelILBasicBlock':
+		return BasicBlockList.__next__(self) # type: ignore
+
+
+class MediumLevelILBasicBlockList(BasicBlockList):
+
+	def __repr__(self):
+		return f"<MediumLevelILBasicBlockList {len(self)} BasicBlocks: {list(self)}>"
+
+	def __getitem__(self, i:Union[int, slice]) -> Union['mediumlevelil.MediumLevelILBasicBlock', List['mediumlevelil.MediumLevelILBasicBlock']]:
+		return BasicBlockList.__getitem__(self, i) # type: ignore
+
+	def __next__(self) -> 'mediumlevelil.MediumLevelILBasicBlock':
+		return BasicBlockList.__next__(self) # type: ignore
+
+
+class HighLevelILBasicBlockList(BasicBlockList):
+
+	def __repr__(self):
+		return f"<HighLevelILBasicBlockList {len(self)} BasicBlocks: {list(self)}>"
+
+	def __getitem__(self, i:Union[int, slice]) -> Union['highlevelil.HighLevelILBasicBlock', List['highlevelil.HighLevelILBasicBlock']]:
+		return BasicBlockList.__getitem__(self, i) # type: ignore
+
+	def __next__(self) -> 'highlevelil.HighLevelILBasicBlock':
+		return BasicBlockList.__next__(self) # type: ignore
+
+
 class TagList:
 	def __init__(self, function:'Function'):
 		self._count = ctypes.c_ulonglong()
diff --git a/python/highlevelil.py b/python/highlevelil.py
index e5b9311d..b8203e75 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -2181,9 +2181,8 @@ class HighLevelILFunction:
 		return HighLevelILBasicBlock(handle, self, self.view)
 
 	@property
-	def basic_blocks(self) -> 'function.BasicBlockList':
-		"""function.BasicBlockList of HighLevelILBasicBlock objects (read-only)"""
-		return function.BasicBlockList(self)
+	def basic_blocks(self) -> 'function.HighLevelILBasicBlockList':
+		return function.HighLevelILBasicBlockList(self)
 
 	@property
 	def instructions(self) -> Generator[HighLevelILInstruction, None, None]:
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 4d320d18..c766b8dd 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -2793,9 +2793,8 @@ class LowLevelILFunction:
 		return LowLevelILBasicBlock(handle, self, self.view)
 
 	@property
-	def basic_blocks(self) -> 'function.BasicBlockList':
-		"""function.BasicBlockList of LowLevelILBasicBlock objects (read-only)"""
-		return function.BasicBlockList(self)
+	def basic_blocks(self) -> 'function.LowLevelILBasicBlockList':
+		return function.LowLevelILBasicBlockList(self)
 
 	@property
 	def instructions(self) -> Generator['LowLevelILInstruction', None, None]:
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index becc22da..38a35de3 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -2693,9 +2693,8 @@ class MediumLevelILFunction:
 		return MediumLevelILBasicBlock(handle, self, self.view)
 
 	@property
-	def basic_blocks(self) -> 'function.BasicBlockList':
-		"""function.BasicBlockList of MediumLevelILBasicBlock objects (read-only)"""
-		return function.BasicBlockList(self)
+	def basic_blocks(self) -> 'function.MediumLevelILBasicBlockList':
+		return function.MediumLevelILBasicBlockList(self)
 
 	@property
 	def instructions(self) -> Generator[MediumLevelILInstruction, None, None]: