Coverage for src / ethereum_types / enum.py: 100%

1""" 

2Support for enumerations. 

3""" 

4 

5from enum import KEEP, Enum, EnumType, Flag, FlagBoundary 

6from typing import Any, Callable, SupportsInt, TypeVar, Union 

7 

8from typing_extensions import assert_never, override 

9 

10from .numeric import Uint 

11 

12_T = TypeVar("_T") 

13 

14 

15def _copy_type(f: _T) -> Callable[[Any], _T]: 

16 # See: https://github.com/python/typing/issues/769#issuecomment-903760354 

17 del f 

18 return lambda x: x 

19 

20 

21class _UintEnumType(EnumType, type): 

22 @_copy_type(EnumType.__call__) 

23 @override 

24 def __call__(cls, value, *args, **kwargs): # type: ignore[no-untyped-def] 

25 # Allow syntax like `MyUintEnum(3)` instead of `MyUintEnum(Uint(3))`. 

26 # Required because we implement the dunder methods in `Unsigned` as 

27 # `Self(int + int)`. 

28 return super().__call__(Uint(value), *args, **kwargs) 

29 

30 

31class UintEnum(Uint, Enum, metaclass=_UintEnumType): 

32 """ 

33 Base class for creating enumerated constants that are also subclasses of 

34 [`Uint`]. 

35 

36 `UintEnum` is similar to [`IntEnum`], but its members are subclasses of 

37 [`Uint`] instead of `int`. 

38 

39 **Note:** Unlike [`IntEnum`], `UintEnum` integer operations that would 

40 result in a numeric value without a corresponding enumeration member raise 

41 a `ValueError` instead of silently converting to a [`Uint`]. 

42 

43 ```python 

44 >>> from ethereum_types.numeric import Uint 

45 >>> from ethereum_types.enum import UintEnum 

46 >>> class Color(UintEnum): 

47 ... RED = Uint(1) 

48 ... BLUE = Uint(2) 

49 ... 

50 >>> Color.RED 

51 <Color.RED: 1> 

52 ``` 

53 

54 [`Uint`]: ref:ethereum_types.numeric.Uint 

55 [`IntEnum`]: https://docs.python.org/3/library/enum.html#enum.IntEnum 

56 """ 

57 

58 __repr__ = Enum.__repr__ 

59 

60 

61class _UintFlagType(_UintEnumType): 

62 _boundary_: FlagBoundary 

63 

64 @_copy_type(EnumType.__call__) 

65 @override 

66 def __call__(cls, value, *args, **kwargs): # type: ignore[no-untyped-def] 

67 result = super().__call__(value, *args, **kwargs) 

68 if isinstance(result, int): 

69 assert cls._boundary_ is not KEEP 

70 value = result 

71 else: 

72 if cls._boundary_ is KEEP or result._name_ in cls.__members__: 

73 return result 

74 value = result._value_ 

75 

76 if cls._boundary_ is FlagBoundary.EJECT: 

77 return Uint(value) 

78 

79 if value == 0: 

80 member: UintFlag = Uint.__new__(cls) # type: ignore[arg-type] 

81 Uint.__init__(member, 0) 

82 member._value_ = 0 

83 member._name_ = "None" 

84 return member 

85 

86 return result 

87 

88 assert_never(cls._boundary_) # pragma: nocover 

89 

90 

91class UintFlag(Uint, Flag, metaclass=_UintFlagType, boundary=KEEP): 

92 def __new__(cls, value: SupportsInt) -> "UintFlag": 

93 member = Uint.__new__(cls) 

94 Uint.__init__(member, value) 

95 member._value_ = int(value) 

96 return member 

97 

98 @classmethod 

99 @override 

100 def _missing_(cls, value: object) -> Union[None, "Uint", int]: 

101 if not isinstance(value, SupportsInt): 

102 return None 

103 

104 int_value = int(value) 

105 missed = super()._missing_(int_value) 

106 if isinstance(missed, UintFlag): 

107 UintFlag.__init__(missed, int_value) 

108 return missed 

109 

110 assert isinstance(missed, (Uint, int)) 

111 return missed 

112 

113 @override 

114 def __repr__(self) -> str: 

115 value = Uint(self._value_) 

116 return f"<{self.__class__.__name__}.{self._name_}: {value!r}>"