_strategies.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236
  1. from typing import (
  2. Callable,
  3. Optional,
  4. Union,
  5. )
  6. from eth_typing.abi import (
  7. TypeStr,
  8. )
  9. from eth_utils import (
  10. to_checksum_address,
  11. )
  12. from hypothesis import (
  13. strategies as st,
  14. )
  15. from eth_abi.grammar import (
  16. ABIType,
  17. BasicType,
  18. TupleType,
  19. normalize,
  20. parse,
  21. )
  22. from eth_abi.registry import (
  23. BaseEquals,
  24. BaseRegistry,
  25. Lookup,
  26. PredicateMapping,
  27. has_arrlist,
  28. is_base_tuple,
  29. )
  30. from eth_abi.utils.numeric import (
  31. scale_places,
  32. )
  33. StrategyFactory = Callable[[ABIType, "StrategyRegistry"], st.SearchStrategy]
  34. StrategyRegistration = Union[st.SearchStrategy, StrategyFactory]
  35. class StrategyRegistry(BaseRegistry):
  36. def __init__(self) -> None:
  37. self._strategies = PredicateMapping("strategy registry")
  38. def register_strategy(
  39. self,
  40. lookup: Lookup,
  41. registration: StrategyRegistration,
  42. label: Optional[str] = None,
  43. ) -> None:
  44. self._register(self._strategies, lookup, registration, label=label)
  45. def unregister_strategy(self, lookup_or_label: Lookup) -> None:
  46. self._unregister(self._strategies, lookup_or_label)
  47. def get_strategy(self, type_str: TypeStr) -> st.SearchStrategy:
  48. """
  49. Returns a hypothesis strategy for the given ABI type.
  50. :param type_str: The canonical string representation of the ABI type
  51. for which a hypothesis strategy should be returned.
  52. :returns: A hypothesis strategy for generating Python values that are
  53. encodable as values of the given ABI type.
  54. """
  55. registration = self._get_registration(self._strategies, type_str)
  56. if isinstance(registration, st.SearchStrategy):
  57. # If a hypothesis strategy was registered, just return it
  58. return registration
  59. else:
  60. # Otherwise, assume the factory is a callable. Call it with the abi
  61. # type to get an appropriate hypothesis strategy.
  62. normalized_type_str = normalize(type_str)
  63. abi_type = parse(normalized_type_str)
  64. strategy = registration(abi_type, self)
  65. return strategy # type: ignore[no-any-return] # clarify return type
  66. def get_uint_strategy(
  67. abi_type: BasicType, registry: StrategyRegistry
  68. ) -> st.SearchStrategy:
  69. bits = abi_type.sub
  70. return st.integers(
  71. min_value=0,
  72. max_value=2**bits - 1,
  73. )
  74. def get_int_strategy(
  75. abi_type: BasicType, registry: StrategyRegistry
  76. ) -> st.SearchStrategy:
  77. bits = abi_type.sub
  78. return st.integers(
  79. min_value=-(2 ** (bits - 1)),
  80. max_value=2 ** (bits - 1) - 1,
  81. )
  82. address_strategy = st.binary(min_size=20, max_size=20).map(to_checksum_address)
  83. bool_strategy = st.booleans()
  84. def get_ufixed_strategy(
  85. abi_type: BasicType, registry: StrategyRegistry
  86. ) -> st.SearchStrategy:
  87. bits, places = abi_type.sub
  88. return st.decimals(
  89. min_value=0,
  90. max_value=2**bits - 1,
  91. places=0,
  92. ).map(scale_places(places))
  93. def get_fixed_strategy(
  94. abi_type: BasicType, registry: StrategyRegistry
  95. ) -> st.SearchStrategy:
  96. bits, places = abi_type.sub
  97. return st.decimals(
  98. min_value=-(2 ** (bits - 1)),
  99. max_value=2 ** (bits - 1) - 1,
  100. places=0,
  101. ).map(scale_places(places))
  102. def get_bytes_strategy(
  103. abi_type: BasicType, registry: StrategyRegistry
  104. ) -> st.SearchStrategy:
  105. num_bytes = abi_type.sub
  106. return st.binary(
  107. min_size=num_bytes,
  108. max_size=num_bytes,
  109. )
  110. bytes_strategy = st.binary(min_size=0, max_size=4096)
  111. string_strategy = st.text()
  112. def get_array_strategy(
  113. abi_type: ABIType, registry: StrategyRegistry
  114. ) -> st.SearchStrategy:
  115. item_type = abi_type.item_type
  116. item_type_str = item_type.to_type_str()
  117. item_strategy = registry.get_strategy(item_type_str)
  118. last_dim = abi_type.arrlist[-1]
  119. if len(last_dim) == 0:
  120. # Is dynamic list. Don't restrict length.
  121. return st.lists(item_strategy)
  122. else:
  123. # Is static list. Restrict length.
  124. dim_size = last_dim[0]
  125. return st.lists(item_strategy, min_size=dim_size, max_size=dim_size)
  126. def get_tuple_strategy(
  127. abi_type: TupleType, registry: StrategyRegistry
  128. ) -> st.SearchStrategy:
  129. component_strategies = [
  130. registry.get_strategy(comp_abi_type.to_type_str())
  131. for comp_abi_type in abi_type.components
  132. ]
  133. return st.tuples(*component_strategies)
  134. strategy_registry = StrategyRegistry()
  135. strategy_registry.register_strategy(
  136. BaseEquals("uint"),
  137. get_uint_strategy,
  138. label="uint",
  139. )
  140. strategy_registry.register_strategy(
  141. BaseEquals("int"),
  142. get_int_strategy,
  143. label="int",
  144. )
  145. strategy_registry.register_strategy(
  146. BaseEquals("address", with_sub=False),
  147. address_strategy,
  148. label="address",
  149. )
  150. strategy_registry.register_strategy(
  151. BaseEquals("bool", with_sub=False),
  152. bool_strategy,
  153. label="bool",
  154. )
  155. strategy_registry.register_strategy(
  156. BaseEquals("ufixed"),
  157. get_ufixed_strategy,
  158. label="ufixed",
  159. )
  160. strategy_registry.register_strategy(
  161. BaseEquals("fixed"),
  162. get_fixed_strategy,
  163. label="fixed",
  164. )
  165. strategy_registry.register_strategy(
  166. BaseEquals("bytes", with_sub=True),
  167. get_bytes_strategy,
  168. label="bytes<M>",
  169. )
  170. strategy_registry.register_strategy(
  171. BaseEquals("bytes", with_sub=False),
  172. bytes_strategy,
  173. label="bytes",
  174. )
  175. strategy_registry.register_strategy(
  176. BaseEquals("function", with_sub=False),
  177. get_bytes_strategy,
  178. label="function",
  179. )
  180. strategy_registry.register_strategy(
  181. BaseEquals("string", with_sub=False),
  182. string_strategy,
  183. label="string",
  184. )
  185. strategy_registry.register_strategy(
  186. has_arrlist,
  187. get_array_strategy,
  188. label="has_arrlist",
  189. )
  190. strategy_registry.register_strategy(
  191. is_base_tuple,
  192. get_tuple_strategy,
  193. label="is_base_tuple",
  194. )
  195. get_abi_strategy = strategy_registry.get_strategy