Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/dissect/cstruct/types/enum.py: 35%

1from __future__ import annotations 

2 

3import sys 

4from enum import Enum as _Enum 

5from enum import EnumMeta, IntEnum, IntFlag 

6from typing import TYPE_CHECKING, Any, BinaryIO, TypeVar, overload 

7 

8from dissect.cstruct.types.base import Array, BaseType, MetaType 

9 

10if TYPE_CHECKING: 

11 from typing_extensions import Self 

12 

13 from dissect.cstruct.cstruct import cstruct 

14 

15 

16PY_311 = sys.version_info >= (3, 11, 0) 

17PY_312 = sys.version_info >= (3, 12, 0) 

18 

19_S = TypeVar("_S") 

20 

21 

22class EnumMetaType(EnumMeta, MetaType): 

23 type: type[BaseType] 

24 

25 @overload 

26 def __call__(cls, value: cstruct, name: str, type_: type[BaseType], *args, **kwargs) -> type[Enum]: ... 

27 

28 @overload 

29 def __call__(cls: type[_S], value: int | BinaryIO | bytes) -> _S: ... 

30 

31 def __call__( 

32 cls, 

33 value: cstruct | int | BinaryIO | bytes | None = None, 

34 name: str | None = None, 

35 type_: type[BaseType] | None = None, 

36 *args, 

37 **kwargs, 

38 ) -> Enum | type[Enum]: 

39 if name is None: 

40 if value is None: 

41 value = cls.type.__default__() 

42 

43 if not isinstance(value, int): 

44 # value is a parsable value 

45 value = cls.type(value) 

46 

47 return super().__call__(value) 

48 

49 # We are constructing a new Enum class 

50 # cs is the cstruct instance, but we can't isinstance check it due to circular imports 

51 cs = value 

52 if not issubclass(type_, int): 

53 raise TypeError("Enum can only be created from int type") 

54 

55 enum_cls = super().__call__(name, *args, **kwargs) 

56 enum_cls.cs = cs 

57 enum_cls.type = type_ 

58 enum_cls.size = type_.size 

59 enum_cls.dynamic = type_.dynamic 

60 enum_cls.alignment = type_.alignment 

61 

62 _fix_alias_members(enum_cls) 

63 

64 return enum_cls 

65 

66 @overload 

67 def __getitem__(cls: type[_S], name: str) -> _S: ... 

68 

69 @overload 

70 def __getitem__(cls: type[_S], name: int) -> Array: ... 

71 

72 def __getitem__(cls: type[_S], name: str | int) -> _S | Array: 

73 if isinstance(name, str): 

74 return super().__getitem__(name) 

75 return MetaType.__getitem__(cls, name) 

76 

77 __len__ = MetaType.__len__ 

78 

79 def __contains__(cls, value: Any) -> bool: 

80 # We used to let stdlib enum handle `__contains__``` but this commit is incompatible with our API: 

81 # https://github.com/python/cpython/commit/8a9aee71268c77867d3cc96d43cbbdcbe8c0e1e8 

82 if isinstance(value, cls): 

83 return True 

84 return value in cls._value2member_map_ 

85 

86 def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self: 

87 return cls(cls.type._read(stream, context)) 

88 

89 def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> list[Self]: 

90 return list(map(cls, cls.type._read_array(stream, count, context))) 

91 

92 def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> list[Self]: 

93 return list(map(cls, cls.type._read_0(stream, context))) 

94 

95 def _write(cls, stream: BinaryIO, data: Enum) -> int: 

96 return cls.type._write(stream, data.value) 

97 

98 def _write_array(cls, stream: BinaryIO, array: list[BaseType | int]) -> int: 

99 data = [entry.value if isinstance(entry, _Enum) else entry for entry in array] 

100 return cls.type._write_array(stream, data) 

101 

102 def _write_0(cls, stream: BinaryIO, array: list[BaseType | int]) -> int: 

103 data = [entry.value if isinstance(entry, _Enum) else entry for entry in array] 

104 return cls._write_array(stream, [*data, cls.type.__default__()]) 

105 

106 

107def _fix_alias_members(cls: type[Enum]) -> None: 

108 # Emulate aenum NoAlias behaviour 

109 # https://github.com/ethanfurman/aenum/blob/master/aenum/doc/aenum.rst 

110 if len(cls._member_names_) == len(cls._member_map_): 

111 return 

112 

113 for name, member in cls._member_map_.items(): 

114 if name != member.name: 

115 new_member = int.__new__(cls, member.value) 

116 new_member._name_ = name 

117 new_member._value_ = member.value 

118 

119 type.__setattr__(cls, name, new_member) 

120 cls._member_names_.append(name) 

121 cls._member_map_[name] = new_member 

122 cls._value2member_map_[member.value] = new_member 

123 

124 

125class Enum(BaseType, IntEnum, metaclass=EnumMetaType): 

126 """Enum type supercharged with cstruct functionality. 

127 

128 Enums are (mostly) compatible with the Python 3 standard library ``IntEnum`` with some notable differences: 

129 - Duplicate members are their own unique member instead of being an alias 

130 - Non-existing values are allowed and handled similarly to ``IntFlag``: ``<Enum: 0>`` 

131 - Enum members are only considered equal if the enum class is the same 

132 

133 Enums can be made using any integer type. 

134 

135 Example: 

136 When using the default C-style parser, the following syntax is supported:: 

137 

138 enum <name> [: <type>] { 

139 <values> 

140 }; 

141 

142 For example, an enum that has A=1, B=5 and C=6 could be written like so:: 

143 

144 enum Test : uint16 { 

145 A, B=5, C 

146 }; 

147 """ 

148 

149 if PY_311: 

150 

151 def __repr__(self) -> str: 

152 # Use the IntFlag repr as a base since it handles unknown values the way we want it 

153 # I.e. <Color: 255> instead of <Color.None: 255> 

154 result = IntFlag.__repr__(self) 

155 if not self.__class__.__name__: 

156 # Deal with anonymous enums by stripping off the first bit 

157 # I.e. <.RED: 1> -> <RED: 1> 

158 result = f"<{result[2:]}" 

159 return result 

160 

161 def __str__(self) -> str: 

162 # We differentiate with standard Python enums in that we use a more descriptive str representation 

163 # Standard Python enums just use the integer value as str, we use EnumName.ValueName 

164 # In case of anonymous enums, we just use the ValueName 

165 # In case of unknown members, we use the integer value (in combination with the EnumName if there is one) 

166 base = f"{self.__class__.__name__}." if self.__class__.__name__ else "" 

167 value = self.name if self.name is not None else str(self.value) 

168 return f"{base}{value}" 

169 

170 else: 

171 

172 def __repr__(self) -> str: 

173 name = self.__class__.__name__ 

174 if self._name_ is not None: 

175 if name: 

176 name += "." 

177 name += self._name_ 

178 return f"<{name}: {self._value_!r}>" 

179 

180 def __str__(self) -> str: 

181 base = f"{self.__class__.__name__}." if self.__class__.__name__ else "" 

182 value = self._name_ if self._name_ is not None else str(self._value_) 

183 return f"{base}{value}" 

184 

185 def __eq__(self, other: int | Enum) -> bool: 

186 if isinstance(other, Enum) and other.__class__ is not self.__class__: 

187 return False 

188 

189 # Python <= 3.10 compatibility 

190 if isinstance(other, Enum): 

191 other = other.value 

192 

193 return self.value == other 

194 

195 def __ne__(self, value: int | Enum) -> bool: 

196 return not self.__eq__(value) 

197 

198 def __hash__(self) -> int: 

199 return hash((self.__class__, self.name, self.value)) 

200 

201 @classmethod 

202 def _missing_(cls, value: int) -> Self: 

203 # Emulate FlagBoundary.KEEP for enum (allow values other than the defined members) 

204 new_member = int.__new__(cls, value) 

205 new_member._name_ = None 

206 new_member._value_ = value 

207 return new_member