datastructures.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. from typing import (
  2. Any,
  3. Dict,
  4. NamedTuple,
  5. SupportsIndex,
  6. Tuple,
  7. Type,
  8. Union,
  9. overload,
  10. )
  11. import warnings
  12. from eth_keys.datatypes import (
  13. Signature,
  14. )
  15. from eth_typing import (
  16. ChecksumAddress,
  17. )
  18. from eth_utils import (
  19. to_checksum_address,
  20. )
  21. from hexbytes import (
  22. HexBytes,
  23. )
  24. from pydantic import (
  25. BaseModel,
  26. ConfigDict,
  27. Field,
  28. field_serializer,
  29. )
  30. from pydantic.alias_generators import (
  31. to_camel,
  32. )
  33. from pydantic.json_schema import (
  34. DEFAULT_REF_TEMPLATE,
  35. GenerateJsonSchema,
  36. JsonSchemaMode,
  37. JsonSchemaValue,
  38. )
  39. from pydantic_core import (
  40. CoreSchema,
  41. PydanticOmit,
  42. )
  43. class SignedTransaction(
  44. NamedTuple(
  45. "SignedTransaction",
  46. [
  47. ("raw_transaction", HexBytes),
  48. ("hash", HexBytes),
  49. ("r", int),
  50. ("s", int),
  51. ("v", int),
  52. ],
  53. )
  54. ):
  55. @overload
  56. def __getitem__(self, index: SupportsIndex) -> Any:
  57. ...
  58. @overload
  59. def __getitem__(self, index: slice) -> Tuple[Any, ...]:
  60. ...
  61. @overload
  62. def __getitem__(self, index: str) -> Any:
  63. ...
  64. def __getitem__(self, index: Union[SupportsIndex, slice, str]) -> Any:
  65. if isinstance(index, (int, slice)):
  66. return super().__getitem__(index)
  67. elif isinstance(index, str):
  68. return getattr(self, index)
  69. else:
  70. raise TypeError("Index must be an integer, slice, or string")
  71. class SignedMessage(
  72. NamedTuple(
  73. "SignedMessage",
  74. [
  75. ("message_hash", HexBytes),
  76. ("r", int),
  77. ("s", int),
  78. ("v", int),
  79. ("signature", HexBytes),
  80. ],
  81. )
  82. ):
  83. @overload
  84. def __getitem__(self, index: SupportsIndex) -> Any:
  85. ...
  86. @overload
  87. def __getitem__(self, index: slice) -> Tuple[Any, ...]:
  88. ...
  89. @overload
  90. def __getitem__(self, index: str) -> Any:
  91. ...
  92. def __getitem__(self, index: Union[SupportsIndex, slice, str]) -> Any:
  93. if isinstance(index, (int, slice)):
  94. return super().__getitem__(index)
  95. elif isinstance(index, str):
  96. return getattr(self, index)
  97. else:
  98. raise TypeError("Index must be an integer, slice, or string")
  99. class OmitJsonSchema(GenerateJsonSchema):
  100. def handle_invalid_for_json_schema(
  101. self, schema: CoreSchema, error_info: str
  102. ) -> JsonSchemaValue:
  103. raise PydanticOmit
  104. class CustomPydanticModel(BaseModel):
  105. """
  106. A base class for eth-account pydantic models that provides custom JSON-RPC
  107. serialization configuration. JSON-RPC serialization is configured to use
  108. camelCase keys and to exclude fields that are marked as ``exclude=True``. To
  109. serialize a model to the expected JSON-RPC format, use
  110. ``model_dump(by_alias=True)``.
  111. """
  112. model_config = ConfigDict(
  113. arbitrary_types_allowed=True,
  114. populate_by_name=True, # populate by snake_case (python) args
  115. alias_generator=to_camel, # serialize by camelCase (json-rpc) keys
  116. )
  117. def recursive_model_dump(self) -> Any:
  118. # TODO: Remove in next major release. This was introduced because of a bug with
  119. # a ``@field_serializer`` decorator not being applied correctly to nested
  120. # models. This is no longer necessary.
  121. warnings.warn(
  122. "recursive_model_dump() is deprecated. Please use "
  123. "model_dump(by_alias=True) instead.",
  124. DeprecationWarning,
  125. stacklevel=2,
  126. )
  127. return self.model_dump(by_alias=True)
  128. @classmethod
  129. def model_json_schema(
  130. cls,
  131. by_alias: bool = True,
  132. ref_template: str = DEFAULT_REF_TEMPLATE,
  133. # default to ``OmitJsonSchema`` to prevent errors from excluded fields
  134. schema_generator: Type[GenerateJsonSchema] = OmitJsonSchema,
  135. mode: JsonSchemaMode = "validation",
  136. ) -> Dict[str, Any]:
  137. """
  138. Omits excluded fields from the JSON schema, preventing errors that would
  139. otherwise be raised by the default schema generator.
  140. """
  141. return super().model_json_schema(
  142. by_alias=by_alias,
  143. ref_template=ref_template,
  144. schema_generator=schema_generator,
  145. mode=mode,
  146. )
  147. class SignedSetCodeAuthorization(CustomPydanticModel):
  148. chain_id: int
  149. address: bytes
  150. nonce: int
  151. y_parity: int
  152. r: int
  153. s: int
  154. signature: Signature = Field(exclude=True)
  155. authorization_hash: HexBytes = Field(exclude=True)
  156. @field_serializer("address")
  157. @classmethod
  158. def serialize_address(cls, value: bytes) -> ChecksumAddress:
  159. return to_checksum_address(value)
  160. @property
  161. def authority(self) -> bytes:
  162. """
  163. Return the address of the authority that signed the authorization.
  164. In order to prevent any potential confusion or mal-intent, the authority is
  165. always derived from the signature and the authorization hash, rather than
  166. statically assigned. This value should be verified against the expected
  167. authority for a signed authorization.
  168. """
  169. return self.signature.recover_public_key_from_msg_hash(
  170. self.authorization_hash
  171. ).to_canonical_address()