Improve type hints for many __getitem__ impls

5 files changed, 70 insertions, 11 deletions

diff --git a/python/binaryview.py b/python/binaryview.py
index 4a160983..d740ce92 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -30,7 +30,7 @@ import inspect
 import os
 import uuid
 from typing import Callable, Generator, Optional, Union, Tuple, List, Mapping, Any, \
-	Iterator, Iterable, KeysView, ItemsView, ValuesView, Dict
+	Iterator, Iterable, KeysView, ItemsView, ValuesView, Dict, overload
 from dataclasses import dataclass
 from enum import IntFlag
 
@@ -1985,6 +1985,12 @@ class FunctionList:
 		self._n += 1
 		return _function.Function(self._view, func)
 
+	@overload
+	def __getitem__(self, i: int) -> '_function.Function': ...
+
+	@overload
+	def __getitem__(self, i: slice) -> List['_function.Function']: ...
+
 	def __getitem__(self, i: Union[int, slice]) -> Union['_function.Function', List['_function.Function']]:
 		if isinstance(i, int):
 			if i < 0:
diff --git a/python/function.py b/python/function.py
index a3d53feb..e08bb6c3 100644
--- a/python/function.py
+++ b/python/function.py
@@ -21,7 +21,7 @@
 
 import ctypes
 import inspect
-from typing import Generator, Optional, List, Tuple, Union, Mapping, Any, Dict
+from typing import Generator, Optional, List, Tuple, Union, Mapping, Any, Dict, overload
 from dataclasses import dataclass
 
 # Binary Ninja components
@@ -211,6 +211,12 @@ class BasicBlockList:
 		self._n += 1
 		return self._function._instantiate_block(block)
 
+	@overload
+	def __getitem__(self, i: int) -> 'basicblock.BasicBlock': ...
+
+	@overload
+	def __getitem__(self, i: slice) -> List['basicblock.BasicBlock']: ...
+
 	def __getitem__(self, i: Union[int, slice]) -> Union['basicblock.BasicBlock', List['basicblock.BasicBlock']]:
 		if isinstance(i, int):
 			if i < 0:
@@ -237,6 +243,12 @@ class LowLevelILBasicBlockList(BasicBlockList):
 	def __repr__(self):
 		return f"<LowLevelILBasicBlockList {len(self)} BasicBlocks: {list(self)}>"
 
+	@overload
+	def __getitem__(self, i: int) -> 'lowlevelil.LowLevelILBasicBlock': ...
+
+	@overload
+	def __getitem__(self, i: slice) -> List['lowlevelil.LowLevelILBasicBlock']: ...
+
 	def __getitem__(
 	    self, i: Union[int, slice]
 	) -> Union['lowlevelil.LowLevelILBasicBlock', List['lowlevelil.LowLevelILBasicBlock']]:
@@ -250,6 +262,12 @@ class MediumLevelILBasicBlockList(BasicBlockList):
 	def __repr__(self):
 		return f"<MediumLevelILBasicBlockList {len(self)} BasicBlocks: {list(self)}>"
 
+	@overload
+	def __getitem__(self, i: int) -> 'mediumlevelil.MediumLevelILBasicBlock': ...
+
+	@overload
+	def __getitem__(self, i: slice) -> List['mediumlevelil.MediumLevelILBasicBlock']: ...
+
 	def __getitem__(
 	    self, i: Union[int, slice]
 	) -> Union['mediumlevelil.MediumLevelILBasicBlock', List['mediumlevelil.MediumLevelILBasicBlock']]:
@@ -263,6 +281,12 @@ class HighLevelILBasicBlockList(BasicBlockList):
 	def __repr__(self):
 		return f"<HighLevelILBasicBlockList {len(self)} BasicBlocks: {list(self)}>"
 
+	@overload
+	def __getitem__(self, i: int) -> 'highlevelil.HighLevelILBasicBlock': ...
+
+	@overload
+	def __getitem__(self, i: slice) -> List['highlevelil.HighLevelILBasicBlock']: ...
+
 	def __getitem__(
 	    self, i: Union[int, slice]
 	) -> Union['highlevelil.HighLevelILBasicBlock', List['highlevelil.HighLevelILBasicBlock']]:
@@ -304,10 +328,15 @@ class TagList:
 		self._n += 1
 		return arch, address, binaryview.Tag(core_tag)
 
+	@overload
+	def __getitem__(self, i: int) -> Tuple['architecture.Architecture', int, 'binaryview.Tag']: ...
+
+	@overload
+	def __getitem__(self, i: slice) -> List[Tuple['architecture.Architecture', int, 'binaryview.Tag']]: ...
+
 	def __getitem__(
 	    self, i: Union[int, slice]
-	) -> Union[Tuple['architecture.Architecture', int, 'binaryview.Tag'], List[Tuple['architecture.Architecture', int,
-	                                                                                 'binaryview.Tag']]]:
+	) -> Union[Tuple['architecture.Architecture', int, 'binaryview.Tag'], List[Tuple['architecture.Architecture', int, 'binaryview.Tag']]]:
 		if isinstance(i, int):
 			if i < 0:
 				i = len(self) + i
@@ -400,7 +429,13 @@ class Function:
 	def __hash__(self):
 		return hash((self.start, self.arch, self.platform))
 
-	def __getitem__(self, i) -> Union['basicblock.BasicBlock', List['basicblock.BasicBlock']]:
+	@overload
+	def __getitem__(self, i: int) -> 'basicblock.BasicBlock': ...
+
+	@overload
+	def __getitem__(self, i: slice) -> List['basicblock.BasicBlock']: ...
+
+	def __getitem__(self, i: Union[int, slice]) -> Union['basicblock.BasicBlock', List['basicblock.BasicBlock']]:
 		return self.basic_blocks[i]
 
 	def __iter__(self) -> Generator['basicblock.BasicBlock', None, None]:
diff --git a/python/highlevelil.py b/python/highlevelil.py
index 2b4d9153..79ac7366 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -20,7 +20,7 @@
 
 import ctypes
 import struct
-from typing import Optional, Generator, List, Union, NewType, Tuple, ClassVar, Mapping, Set, Callable, Any, Iterator
+from typing import Optional, Generator, List, Union, NewType, Tuple, ClassVar, Mapping, Set, Callable, Any, Iterator, overload
 from dataclasses import dataclass
 from enum import Enum
 
@@ -3089,7 +3089,13 @@ class HighLevelILBasicBlock(basicblock.BasicBlock):
 		for idx in range(self.start, self.end):
 			yield self.il_function[idx]
 
-	def __getitem__(self, idx) -> Union[List[HighLevelILInstruction], HighLevelILInstruction]:
+	@overload
+	def __getitem__(self, idx: int) -> 'HighLevelILInstruction': ...
+
+	@overload
+	def __getitem__(self, idx: slice) -> List['HighLevelILInstruction']: ...
+
+	def __getitem__(self, idx: Union[int, slice]) -> Union[List[HighLevelILInstruction], HighLevelILInstruction]:
 		size = self.end - self.start
 		if isinstance(idx, slice):
 			return [self[index] for index in range(*idx.indices(size))]  # type: ignore
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 02586113..e3c3ccab 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -20,7 +20,7 @@
 
 import ctypes
 import struct
-from typing import Generator, List, Optional, Dict, Union, Tuple, NewType, ClassVar, Set, Callable, Any, Iterator
+from typing import Generator, List, Optional, Dict, Union, Tuple, NewType, ClassVar, Set, Callable, Any, Iterator, overload
 from dataclasses import dataclass
 
 # Binary Ninja components
@@ -5503,7 +5503,13 @@ class LowLevelILBasicBlock(basicblock.BasicBlock):
 		for idx in range(self.start, self.end):
 			yield self._il_function[idx]
 
-	def __getitem__(self, idx):
+	@overload
+	def __getitem__(self, idx: int) -> 'LowLevelILInstruction': ...
+
+	@overload
+	def __getitem__(self, idx: slice) -> List['LowLevelILInstruction']: ...
+
+	def __getitem__(self, idx: Union[int, slice]) -> Union['LowLevelILInstruction', List['LowLevelILInstruction']]:
 		size = self.end - self.start
 		if isinstance(idx, slice):
 			return [self[index] for index in range(*idx.indices(size))]
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index e857920f..291da4b9 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -21,7 +21,7 @@
 import ctypes
 import struct
 from typing import (Optional, List, Union, Mapping,
-	Generator, NewType, Tuple, ClassVar, Dict, Set, Callable, Any, Iterator)
+	Generator, NewType, Tuple, ClassVar, Dict, Set, Callable, Any, Iterator, overload)
 from dataclasses import dataclass
 from . import deprecation
 
@@ -3957,7 +3957,13 @@ class MediumLevelILBasicBlock(basicblock.BasicBlock):
 		for idx in range(self.start, self.end):
 			yield self._il_function[idx]
 
-	def __getitem__(self, idx) -> Union[List['MediumLevelILInstruction'], 'MediumLevelILInstruction']:
+	@overload
+	def __getitem__(self, idx: int) -> 'MediumLevelILInstruction': ...
+
+	@overload
+	def __getitem__(self, idx: slice) -> List['MediumLevelILInstruction']: ...
+
+	def __getitem__(self, idx: Union[int, slice]) -> Union[List['MediumLevelILInstruction'], 'MediumLevelILInstruction']:
 		size = self.end - self.start
 		if isinstance(idx, slice):
 			return [self[index] for index in range(*idx.indices(size))]  # type: ignore