Merge pull request #36 from akiym/fix-function-signature

Fix function type signatures

ptrlib/binary/packing/chunks.py

@@ -4,12 +4,12 @@
 _U = TypeVar("_U")
 
 @overload
-def chunks(data: _T, size: int, padding: Optional[_T], map: None) -> List[_T]: ...
+def chunks(data: _T, size: int, padding: Optional[_T]=None, map: None=None) -> List[_T]: ...
 
 @overload
-def chunks(data: _T, size: int, padding: Optional[_U], map: Callable[[_T], _U]) -> List[_U]: ...
+def chunks(data: _T, size: int, padding: None=None, map: Optional[Callable[[_T], _U]]=None) -> List[_U]: ...
 
-def chunks(data: _T, size: int, padding: Optional[Union[_T, _U]]=None, map: Optional[Callable[[_T], _U]]=None) -> Union[List[_T], List[_U]]:
+def chunks(data: _T, size: int, padding: Optional[_T]=None, map: Optional[Callable[[_T], _U]]=None) -> Union[List[_T], List[_U]]:
     """Split data into chunks
     Args:
         data      : The target data to split

ptrlib/binary/packing/flat.py

@@ -4,12 +4,12 @@
 _U = TypeVar("_U")
 
 @overload
-def flat(chunks: List[List[_T]], map: None=None) -> List[_T]: ...
+def flat(chunks: List[_T], map: None=None) -> _T: ...
 
 @overload
-def flat(chunks: List[_T], map: Callable[[_T], List[_U]]) -> List[_U]: ...
+def flat(chunks: List[_T], map: Optional[Callable[[_T], _U]]=None) -> _U: ...
 
-def flat(chunks: Union[List[List[_T]], List[_T]], map: Optional[Callable[[_T], List[_U]]]=None) -> Union[List[_T], List[_U]]:
+def flat(chunks: List[_T], map: Optional[Callable[[_T], _U]]=None) -> Union[_T, _U]:
     """Concatnate chunks into a data
     Aimed for the use of crafting ROP chains
 

ptrlib/binary/packing/unpack.py

@@ -1,7 +1,7 @@
 import builtins
 from logging import getLogger
 import struct
-from typing import Type, TypeVar, Union, overload
+from typing import Type, TypeVar, Union
 try:
     from typing import Literal
 except:
@@ -10,8 +10,6 @@
 
 logger = getLogger(__name__)
 
-_T = TypeVar("_T", int, float)
-
 def u8(data: Union[str, bytes], signed: bool=False) -> int:
     if isinstance(data, str):
         data = str2bytes(data)
@@ -30,13 +28,7 @@ def u16(data: Union[str, bytes], byteorder: Literal["little", "big"]='little', s
 
     return int.from_bytes(data, byteorder=byteorder, signed=signed)
 
-@overload
-def u32(data: Union[str, bytes], byteorder: Literal["little", "big"]="little", signed: bool=False, result_type: Type[int]=int) -> int: ...
-
-@overload
-def u32(data: Union[str, bytes], byteorder: Literal["little", "big"]="little", signed: bool=False, result_type: Type[float]=float) -> float: ...
-
-def u32(data: Union[str, bytes], byteorder: Literal["little", "big"]='little', signed: bool=False, result_type: Type[_T]=int) -> _T:
+def u32(data: Union[str, bytes], byteorder: Literal["little", "big"]='little', signed: bool=False, result_type: Union[Type[int], Type[float]]=int) -> int:
     if isinstance(data, str):
         data = str2bytes(data)
 
@@ -52,7 +44,7 @@ def u32(data: Union[str, bytes], byteorder: Literal["little", "big"]='little', s
 
     return int.from_bytes(data, byteorder=byteorder, signed=signed)
 
-def u32f(data: Union[str, bytes], byteorder: Literal["little", "big"]="little"):
+def u32f(data: Union[str, bytes], byteorder: Literal["little", "big"]="little") -> float:
     if isinstance(data, str):
         data = str2bytes(data)
 
@@ -64,13 +56,7 @@ def u32f(data: Union[str, bytes], byteorder: Literal["little", "big"]="little"):
         data
     )[0]
 
-@overload
-def u64(data: Union[str, bytes], byteorder: Literal["little", "big"]="little", signed: bool=False, type: Type[int]=int) -> int: ...
-
-@overload
-def u64(data: Union[str, bytes], byteorder: Literal["little", "big"]="little", signed: bool=False, type: Type[float]=float) -> float: ...
-
-def u64(data: Union[str, bytes], byteorder: Literal["little", "big"]='little', signed: bool=False, type: Type[_T]=int) -> _T:
+def u64(data: Union[str, bytes], byteorder: Literal["little", "big"]='little', signed: bool=False, type: Union[Type[int], Type[float]]=int) -> int:
     if isinstance(data, str):
         data = str2bytes(data)
 
@@ -86,7 +72,7 @@ def u64(data: Union[str, bytes], byteorder: Literal["little", "big"]='little', s
 
     return int.from_bytes(data, byteorder=byteorder, signed=signed)
 
-def u64f(data: Union[str, bytes], byteorder: Literal["little", "big"]="little"):
+def u64f(data: Union[str, bytes], byteorder: Literal["little", "big"]="little") -> float:
     if isinstance(data, str):
         data = str2bytes(data)