Skip to content

Commit 65843d0

Browse files
marci07iqmarci07iq
authored
Type annotation bugfixes (#4740)
* Fix type of new Packet, annotate enums * Revent Packet, fix Enum annotation * Factor enum type into alias
1 parent 267f99a commit 65843d0

2 files changed

Lines changed: 68 additions & 27 deletions

File tree

‎scapy/arch/linux/rtnetlink.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -394,7 +394,7 @@ def default_payload_class(self, payload: bytes) -> Type[Packet]:
394394

395395
class ifinfomsg(Packet):
396396
fields_desc = [
397-
ByteEnumField("ifi_family", 0, socket.AddressFamily), # type: ignore
397+
ByteEnumField("ifi_family", 0, socket.AddressFamily),
398398
ByteField("res", 0),
399399
Field("ifi_type", 0, fmt="=H"),
400400
Field("ifi_index", 0, fmt="=i"),
@@ -483,7 +483,7 @@ def default_payload_class(self, payload: bytes) -> Type[Packet]:
483483

484484
class ifaddrmsg(Packet):
485485
fields_desc = [
486-
ByteEnumField("ifa_family", 0, socket.AddressFamily), # type: ignore
486+
ByteEnumField("ifa_family", 0, socket.AddressFamily),
487487
ByteField("ifa_prefixlen", 0),
488488
FlagsField(
489489
"ifa_flags",
@@ -607,7 +607,7 @@ def default_payload_class(self, payload: bytes) -> Type[Packet]:
607607

608608
class rtmsg(Packet):
609609
fields_desc = [
610-
ByteEnumField("rtm_family", 0, socket.AddressFamily), # type: ignore
610+
ByteEnumField("rtm_family", 0, socket.AddressFamily),
611611
ByteField("rtm_dst_len", 0),
612612
ByteField("rtm_src_len", 0),
613613
ByteField("rtm_tos", 0),

‎scapy/fields.py‎

Lines changed: 65 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1637,7 +1637,7 @@ def __init__(
16371637
pkt_cls=None, # type: Optional[Union[Callable[[bytes], Packet], Type[Packet]]] # noqa: E501
16381638
count_from=None, # type: Optional[Callable[[Packet], int]]
16391639
length_from=None, # type: Optional[Callable[[Packet], int]]
1640-
next_cls_cb=None, # type: Optional[Callable[[Packet, List[BasePacket], Optional[Packet], bytes], Type[Packet]]] # noqa: E501
1640+
next_cls_cb=None, # type: Optional[Callable[[Packet, List[BasePacket], Optional[Packet], bytes], Optional[Type[Packet]]]] # noqa: E501
16411641
max_count=None, # type: Optional[int]
16421642
):
16431643
# type: (...) -> None
@@ -2514,11 +2514,14 @@ def i2repr(self, pkt, x):
25142514
return lhex(self.i2h(pkt, x))
25152515

25162516

2517+
_EnumType = Union[Dict[I, str], Dict[str, I], List[str], DADict[I, str], Type[Enum], Tuple[Callable[[I], str], Callable[[str], I]]] # noqa: E501
2518+
2519+
25172520
class _EnumField(Field[Union[List[I], I], I]):
25182521
def __init__(self,
25192522
name, # type: str
25202523
default, # type: Optional[I]
2521-
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
2524+
enum, # type: _EnumType[I]
25222525
fmt="H", # type: str
25232526
):
25242527
# type: (...) -> None
@@ -2647,7 +2650,7 @@ class CharEnumField(EnumField[str]):
26472650
def __init__(self,
26482651
name, # type: str
26492652
default, # type: str
2650-
enum, # type: Union[Dict[str, str], Tuple[Callable[[str], str], Callable[[str], str]]] # noqa: E501
2653+
enum, # type: _EnumType[str]
26512654
fmt="1s", # type: str
26522655
):
26532656
# type: (...) -> None
@@ -2670,8 +2673,14 @@ def any2i_one(self, pkt, x):
26702673
class BitEnumField(_BitField[Union[List[int], int]], _EnumField[int]):
26712674
__slots__ = EnumField.__slots__
26722675

2673-
def __init__(self, name, default, size, enum, **kwargs):
2674-
# type: (str, Optional[int], int, Dict[int, str], **Any) -> None
2676+
def __init__(self,
2677+
name, # type: str
2678+
default, # type: Optional[int]
2679+
size, # type: int
2680+
enum, # type: _EnumType[int]
2681+
**kwargs # type: Any
2682+
):
2683+
# type: (...) -> None
26752684
_EnumField.__init__(self, name, default, enum)
26762685
_BitField.__init__(self, name, default, size, **kwargs)
26772686

@@ -2694,7 +2703,7 @@ def __init__(self,
26942703
name, # type: str
26952704
default, # type: Optional[int]
26962705
length_from, # type: Callable[[Packet], int]
2697-
enum, # type: Dict[int, str]
2706+
enum, # type: _EnumType[int]
26982707
**kwargs, # type: Any
26992708
):
27002709
# type: (...) -> None
@@ -2718,34 +2727,50 @@ class ShortEnumField(EnumField[int]):
27182727

27192728
def __init__(self,
27202729
name, # type: str
2721-
default, # type: int
2722-
enum, # type: Union[Dict[int, str], Dict[str, int], Tuple[Callable[[int], str], Callable[[str], int]], DADict[int, str]] # noqa: E501
2730+
default, # type: Optional[int]
2731+
enum, # type: _EnumType[int]
27232732
):
27242733
# type: (...) -> None
27252734
super(ShortEnumField, self).__init__(name, default, enum, "H")
27262735

27272736

27282737
class LEShortEnumField(EnumField[int]):
2729-
def __init__(self, name, default, enum):
2730-
# type: (str, int, Union[Dict[int, str], List[str]]) -> None
2738+
def __init__(self,
2739+
name, # type: str
2740+
default, # type: Optional[int]
2741+
enum, # type: _EnumType[int]
2742+
):
2743+
# type: (...) -> None
27312744
super(LEShortEnumField, self).__init__(name, default, enum, "<H")
27322745

27332746

27342747
class LongEnumField(EnumField[int]):
2735-
def __init__(self, name, default, enum):
2736-
# type: (str, int, Union[Dict[int, str], List[str]]) -> None
2748+
def __init__(self,
2749+
name, # type: str
2750+
default, # type: Optional[int]
2751+
enum, # type: _EnumType[int]
2752+
):
2753+
# type: (...) -> None
27372754
super(LongEnumField, self).__init__(name, default, enum, "Q")
27382755

27392756

27402757
class LELongEnumField(EnumField[int]):
2741-
def __init__(self, name, default, enum):
2742-
# type: (str, int, Union[Dict[int, str], List[str]]) -> None
2758+
def __init__(self,
2759+
name, # type: str
2760+
default, # type: Optional[int]
2761+
enum, # type: _EnumType[int]
2762+
):
2763+
# type: (...) -> None
27432764
super(LELongEnumField, self).__init__(name, default, enum, "<Q")
27442765

27452766

27462767
class ByteEnumField(EnumField[int]):
2747-
def __init__(self, name, default, enum):
2748-
# type: (str, Optional[int], Dict[int, str]) -> None
2768+
def __init__(self,
2769+
name, # type: str
2770+
default, # type: Optional[int]
2771+
enum, # type: _EnumType[int]
2772+
):
2773+
# type: (...) -> None
27492774
super(ByteEnumField, self).__init__(name, default, enum, "B")
27502775

27512776

@@ -2766,20 +2791,32 @@ def i2repr_one(self, pkt, x):
27662791

27672792

27682793
class IntEnumField(EnumField[int]):
2769-
def __init__(self, name, default, enum):
2770-
# type: (str, Optional[int], Dict[int, str]) -> None
2794+
def __init__(self,
2795+
name, # type: str
2796+
default, # type: Optional[int]
2797+
enum, # type: _EnumType[int]
2798+
):
2799+
# type: (...) -> None
27712800
super(IntEnumField, self).__init__(name, default, enum, "I")
27722801

27732802

27742803
class SignedIntEnumField(EnumField[int]):
2775-
def __init__(self, name, default, enum):
2776-
# type: (str, Optional[int], Dict[int, str]) -> None
2804+
def __init__(self,
2805+
name, # type: str
2806+
default, # type: Optional[int]
2807+
enum, # type: _EnumType[int]
2808+
):
2809+
# type: (...) -> None
27772810
super(SignedIntEnumField, self).__init__(name, default, enum, "i")
27782811

27792812

27802813
class LEIntEnumField(EnumField[int]):
2781-
def __init__(self, name, default, enum):
2782-
# type: (str, int, Dict[int, str]) -> None
2814+
def __init__(self,
2815+
name, # type: str
2816+
default, # type: Optional[int]
2817+
enum, # type: _EnumType[int]
2818+
):
2819+
# type: (...) -> None
27832820
super(LEIntEnumField, self).__init__(name, default, enum, "<I")
27842821

27852822

@@ -2798,8 +2835,12 @@ def _i2repr(self, pkt, x):
27982835
class LE3BytesEnumField(LEThreeBytesField, _EnumField[int]):
27992836
__slots__ = EnumField.__slots__
28002837

2801-
def __init__(self, name, default, enum):
2802-
# type: (str, Optional[int], Dict[int, str]) -> None
2838+
def __init__(self,
2839+
name, # type: str
2840+
default, # type: Optional[int]
2841+
enum, # type: _EnumType[int]
2842+
):
2843+
# type: (...) -> None
28032844
_EnumField.__init__(self, name, default, enum)
28042845
LEThreeBytesField.__init__(self, name, default)
28052846

0 commit comments

Comments
 (0)