From 72cb4b08ac8f7c9261867cd6c0b18c054d73b96e Mon Sep 17 00:00:00 2001 From: Josh Ferrell Date: Sun, 21 Apr 2024 22:05:45 -0400 Subject: Improve type hints for many __getitem__ impls --- python/highlevelil.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) (limited to 'python/highlevelil.py') 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 -- cgit v1.3.1