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

1""" 

2Dataclass extension that supports immutability. 

3""" 

4 

5from dataclasses import is_dataclass, replace 

6from functools import wraps 

7from typing import ( 

8 Any, 

9 Callable, 

10 Concatenate, 

11 ParamSpec, 

12 Protocol, 

13 TypeVar, 

14 cast, 

15 runtime_checkable, 

16) 

17 

18from mypy_extensions import mypyc_attr 

19 

20 

21@mypyc_attr(native_class=False) 

22@runtime_checkable 

23class SlottedFreezable(Protocol): 

24 """ 

25 A [`Protocol`] implemented by data classes annotated with 

26 [`@slotted_freezable`]. 

27 

28 [`@slotted_freezable`]: ref:ethereum_types.frozen.slotted_freezable 

29 [`Protocol`]: https://docs.python.org/library/typing.html#typing.Protocol 

30 """ 

31 

32 _frozen: bool 

33 

34 

35_MUTATE_MESSAGE = "Mutating frozen dataclasses is not allowed." 

36 

37 

38def _setattr_function(self: Any, attr: str, value: Any) -> None: 

39 if getattr(self, "_frozen", None): 

40 raise AttributeError(_MUTATE_MESSAGE) 

41 else: 

42 object.__setattr__(self, attr, value) 

43 

44 

45def _delattr_function(self: Any, attr: str) -> None: 

46 if self._frozen: 

47 raise AttributeError(_MUTATE_MESSAGE) 

48 else: 

49 object.__delattr__(self, attr) 

50 

51 

52_S = TypeVar("_S", bound=SlottedFreezable) 

53_P = ParamSpec("_P") 

54 

55 

56def _make_init_function( 

57 f: Callable[Concatenate[_S, _P], None], 

58) -> Callable[Concatenate[_S, _P], None]: 

59 @wraps(f) 

60 def init_function(self: _S, *args: _P.args, **kwargs: _P.kwargs) -> None: 

61 will_be_frozen = kwargs.pop("_frozen", True) 

62 assert isinstance(will_be_frozen, bool) 

63 object.__setattr__(self, "_frozen", False) 

64 f(self, *args, **kwargs) 

65 self._frozen = will_be_frozen 

66 

67 return cast("Callable[Concatenate[_S, _P], None]", init_function) 

68 

69 

70def slotted_freezable(cls: Any) -> Any: 

71 """ 

72 Monkey patches a dataclass so it can be frozen by setting `_frozen` to 

73 `True` and uses `__slots__` for efficiency. 

74 

75 Instances will be created frozen by default unless you pass `_frozen=False` 

76 to `__init__`. 

77 """ 

78 cls.__slots__ = ("_frozen", *tuple(cls.__annotations__)) 

79 cls.__init__ = _make_init_function(cls.__init__) 

80 cls.__setattr__ = _setattr_function 

81 cls.__delattr__ = _delattr_function 

82 return type(cls)(cls.__name__, cls.__bases__, dict(cls.__dict__)) 

83 

84 

85S = TypeVar("S") 

86 

87 

88def modify(obj: S, f: Callable[[S], None]) -> S: 

89 """ 

90 Create a copy of `obj` (which must be [`@slotted_freezable`]), and modify 

91 it by applying `f`. The returned copy will be frozen. 

92 

93 [`@slotted_freezable`]: ref:ethereum_types.frozen.slotted_freezable 

94 """ 

95 assert is_dataclass(obj) 

96 assert isinstance(obj, SlottedFreezable) 

97 new_obj = replace(obj, _frozen=False) # type: ignore[unreachable] 

98 f(new_obj) 

99 new_obj._frozen = True 

100 return new_obj