From 72cb4b08ac8f7c9261867cd6c0b18c054d73b96e Mon Sep 17 00:00:00 2001 From: Josh Ferrell Date: Sun, 21 Apr 2024 22:05:45 -0400 Subject: [PATCH] Improve type hints for many __getitem__ impls --- python/binaryview.py | 8 +++++++- python/function.py | 43 +++++++++++++++++++++++++++++++++++++---- python/highlevelil.py | 10 ++++++++-- python/lowlevelil.py | 10 ++++++++-- python/mediumlevelil.py | 10 ++++++++-- 5 files changed, 70 insertions(+), 11 deletions(-) diff --git a/python/binaryview.py b/python/binaryview.py index 4a1609836..d740ce924 100644 --- a/python/binaryview.py +++ b/python/binaryview.py @@ -30,7 +30,7 @@ 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 @@ def __next__(self): 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 a3d53febb..e08bb6c3d 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 @@ def __next__(self) -> 'basicblock.BasicBlock': 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"" + @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"" + @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"" + @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 @@ def __next__(self) -> Tuple['architecture.Architecture', int, 'binaryview.Tag']: 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 @@ def __ge__(self, other: 'Function') -> bool: 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 2b4d9153b..79ac73668 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 @@ def __iter__(self) -> Generator[HighLevelILInstruction, None, None]: 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 025861139..e3c3ccab8 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 @@ def __iter__(self) -> Generator['LowLevelILInstruction', None, None]: 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 e857920f9..291da4b9d 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 @@ def __iter__(self): 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