diff --git a/scapy/arch/linux/rtnetlink.py b/scapy/arch/linux/rtnetlink.py index e57d2459c8c..d5d267df10a 100644 --- a/scapy/arch/linux/rtnetlink.py +++ b/scapy/arch/linux/rtnetlink.py @@ -394,7 +394,7 @@ def default_payload_class(self, payload: bytes) -> Type[Packet]: class ifinfomsg(Packet): fields_desc = [ - ByteEnumField("ifi_family", 0, socket.AddressFamily), # type: ignore + ByteEnumField("ifi_family", 0, socket.AddressFamily), ByteField("res", 0), Field("ifi_type", 0, fmt="=H"), Field("ifi_index", 0, fmt="=i"), @@ -483,7 +483,7 @@ def default_payload_class(self, payload: bytes) -> Type[Packet]: class ifaddrmsg(Packet): fields_desc = [ - ByteEnumField("ifa_family", 0, socket.AddressFamily), # type: ignore + ByteEnumField("ifa_family", 0, socket.AddressFamily), ByteField("ifa_prefixlen", 0), FlagsField( "ifa_flags", @@ -607,7 +607,7 @@ def default_payload_class(self, payload: bytes) -> Type[Packet]: class rtmsg(Packet): fields_desc = [ - ByteEnumField("rtm_family", 0, socket.AddressFamily), # type: ignore + ByteEnumField("rtm_family", 0, socket.AddressFamily), ByteField("rtm_dst_len", 0), ByteField("rtm_src_len", 0), ByteField("rtm_tos", 0), diff --git a/scapy/fields.py b/scapy/fields.py index 443a2d13124..ee513173724 100644 --- a/scapy/fields.py +++ b/scapy/fields.py @@ -1637,7 +1637,7 @@ def __init__( pkt_cls=None, # type: Optional[Union[Callable[[bytes], Packet], Type[Packet]]] # noqa: E501 count_from=None, # type: Optional[Callable[[Packet], int]] length_from=None, # type: Optional[Callable[[Packet], int]] - next_cls_cb=None, # type: Optional[Callable[[Packet, List[BasePacket], Optional[Packet], bytes], Type[Packet]]] # noqa: E501 + next_cls_cb=None, # type: Optional[Callable[[Packet, List[BasePacket], Optional[Packet], bytes], Optional[Type[Packet]]]] # noqa: E501 max_count=None, # type: Optional[int] ): # type: (...) -> None @@ -2514,11 +2514,14 @@ def i2repr(self, pkt, x): return lhex(self.i2h(pkt, x)) +_EnumType = Union[Dict[I, str], Dict[str, I], List[str], DADict[I, str], Type[Enum], Tuple[Callable[[I], str], Callable[[str], I]]] # noqa: E501 + + class _EnumField(Field[Union[List[I], I], I]): def __init__(self, name, # type: str default, # type: Optional[I] - enum, # type: Union[Dict[I, str], Dict[str, I], List[str], DADict[I, str], Type[Enum], Tuple[Callable[[I], str], Callable[[str], I]]] # noqa: E501 + enum, # type: _EnumType[I] fmt="H", # type: str ): # type: (...) -> None @@ -2647,7 +2650,7 @@ class CharEnumField(EnumField[str]): def __init__(self, name, # type: str default, # type: str - enum, # type: Union[Dict[str, str], Tuple[Callable[[str], str], Callable[[str], str]]] # noqa: E501 + enum, # type: _EnumType[str] fmt="1s", # type: str ): # type: (...) -> None @@ -2670,8 +2673,14 @@ def any2i_one(self, pkt, x): class BitEnumField(_BitField[Union[List[int], int]], _EnumField[int]): __slots__ = EnumField.__slots__ - def __init__(self, name, default, size, enum, **kwargs): - # type: (str, Optional[int], int, Dict[int, str], **Any) -> None + def __init__(self, + name, # type: str + default, # type: Optional[int] + size, # type: int + enum, # type: _EnumType[int] + **kwargs # type: Any + ): + # type: (...) -> None _EnumField.__init__(self, name, default, enum) _BitField.__init__(self, name, default, size, **kwargs) @@ -2694,7 +2703,7 @@ def __init__(self, name, # type: str default, # type: Optional[int] length_from, # type: Callable[[Packet], int] - enum, # type: Dict[int, str] + enum, # type: _EnumType[int] **kwargs, # type: Any ): # type: (...) -> None @@ -2718,34 +2727,50 @@ class ShortEnumField(EnumField[int]): def __init__(self, name, # type: str - default, # type: int - enum, # type: Union[Dict[int, str], Dict[str, int], Tuple[Callable[[int], str], Callable[[str], int]], DADict[int, str]] # noqa: E501 + default, # type: Optional[int] + enum, # type: _EnumType[int] ): # type: (...) -> None super(ShortEnumField, self).__init__(name, default, enum, "H") class LEShortEnumField(EnumField[int]): - def __init__(self, name, default, enum): - # type: (str, int, Union[Dict[int, str], List[str]]) -> None + def __init__(self, + name, # type: str + default, # type: Optional[int] + enum, # type: _EnumType[int] + ): + # type: (...) -> None super(LEShortEnumField, self).__init__(name, default, enum, " None + def __init__(self, + name, # type: str + default, # type: Optional[int] + enum, # type: _EnumType[int] + ): + # type: (...) -> None super(LongEnumField, self).__init__(name, default, enum, "Q") class LELongEnumField(EnumField[int]): - def __init__(self, name, default, enum): - # type: (str, int, Union[Dict[int, str], List[str]]) -> None + def __init__(self, + name, # type: str + default, # type: Optional[int] + enum, # type: _EnumType[int] + ): + # type: (...) -> None super(LELongEnumField, self).__init__(name, default, enum, " None + def __init__(self, + name, # type: str + default, # type: Optional[int] + enum, # type: _EnumType[int] + ): + # type: (...) -> None super(ByteEnumField, self).__init__(name, default, enum, "B") @@ -2766,20 +2791,32 @@ def i2repr_one(self, pkt, x): class IntEnumField(EnumField[int]): - def __init__(self, name, default, enum): - # type: (str, Optional[int], Dict[int, str]) -> None + def __init__(self, + name, # type: str + default, # type: Optional[int] + enum, # type: _EnumType[int] + ): + # type: (...) -> None super(IntEnumField, self).__init__(name, default, enum, "I") class SignedIntEnumField(EnumField[int]): - def __init__(self, name, default, enum): - # type: (str, Optional[int], Dict[int, str]) -> None + def __init__(self, + name, # type: str + default, # type: Optional[int] + enum, # type: _EnumType[int] + ): + # type: (...) -> None super(SignedIntEnumField, self).__init__(name, default, enum, "i") class LEIntEnumField(EnumField[int]): - def __init__(self, name, default, enum): - # type: (str, int, Dict[int, str]) -> None + def __init__(self, + name, # type: str + default, # type: Optional[int] + enum, # type: _EnumType[int] + ): + # type: (...) -> None super(LEIntEnumField, self).__init__(name, default, enum, " None + def __init__(self, + name, # type: str + default, # type: Optional[int] + enum, # type: _EnumType[int] + ): + # type: (...) -> None _EnumField.__init__(self, name, default, enum) LEThreeBytesField.__init__(self, name, default)