| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236 |
- from typing import (
- Callable,
- Optional,
- Union,
- )
- from eth_typing.abi import (
- TypeStr,
- )
- from eth_utils import (
- to_checksum_address,
- )
- from hypothesis import (
- strategies as st,
- )
- from eth_abi.grammar import (
- ABIType,
- BasicType,
- TupleType,
- normalize,
- parse,
- )
- from eth_abi.registry import (
- BaseEquals,
- BaseRegistry,
- Lookup,
- PredicateMapping,
- has_arrlist,
- is_base_tuple,
- )
- from eth_abi.utils.numeric import (
- scale_places,
- )
- StrategyFactory = Callable[[ABIType, "StrategyRegistry"], st.SearchStrategy]
- StrategyRegistration = Union[st.SearchStrategy, StrategyFactory]
- class StrategyRegistry(BaseRegistry):
- def __init__(self) -> None:
- self._strategies = PredicateMapping("strategy registry")
- def register_strategy(
- self,
- lookup: Lookup,
- registration: StrategyRegistration,
- label: Optional[str] = None,
- ) -> None:
- self._register(self._strategies, lookup, registration, label=label)
- def unregister_strategy(self, lookup_or_label: Lookup) -> None:
- self._unregister(self._strategies, lookup_or_label)
- def get_strategy(self, type_str: TypeStr) -> st.SearchStrategy:
- """
- Returns a hypothesis strategy for the given ABI type.
- :param type_str: The canonical string representation of the ABI type
- for which a hypothesis strategy should be returned.
- :returns: A hypothesis strategy for generating Python values that are
- encodable as values of the given ABI type.
- """
- registration = self._get_registration(self._strategies, type_str)
- if isinstance(registration, st.SearchStrategy):
- # If a hypothesis strategy was registered, just return it
- return registration
- else:
- # Otherwise, assume the factory is a callable. Call it with the abi
- # type to get an appropriate hypothesis strategy.
- normalized_type_str = normalize(type_str)
- abi_type = parse(normalized_type_str)
- strategy = registration(abi_type, self)
- return strategy # type: ignore[no-any-return] # clarify return type
- def get_uint_strategy(
- abi_type: BasicType, registry: StrategyRegistry
- ) -> st.SearchStrategy:
- bits = abi_type.sub
- return st.integers(
- min_value=0,
- max_value=2**bits - 1,
- )
- def get_int_strategy(
- abi_type: BasicType, registry: StrategyRegistry
- ) -> st.SearchStrategy:
- bits = abi_type.sub
- return st.integers(
- min_value=-(2 ** (bits - 1)),
- max_value=2 ** (bits - 1) - 1,
- )
- address_strategy = st.binary(min_size=20, max_size=20).map(to_checksum_address)
- bool_strategy = st.booleans()
- def get_ufixed_strategy(
- abi_type: BasicType, registry: StrategyRegistry
- ) -> st.SearchStrategy:
- bits, places = abi_type.sub
- return st.decimals(
- min_value=0,
- max_value=2**bits - 1,
- places=0,
- ).map(scale_places(places))
- def get_fixed_strategy(
- abi_type: BasicType, registry: StrategyRegistry
- ) -> st.SearchStrategy:
- bits, places = abi_type.sub
- return st.decimals(
- min_value=-(2 ** (bits - 1)),
- max_value=2 ** (bits - 1) - 1,
- places=0,
- ).map(scale_places(places))
- def get_bytes_strategy(
- abi_type: BasicType, registry: StrategyRegistry
- ) -> st.SearchStrategy:
- num_bytes = abi_type.sub
- return st.binary(
- min_size=num_bytes,
- max_size=num_bytes,
- )
- bytes_strategy = st.binary(min_size=0, max_size=4096)
- string_strategy = st.text()
- def get_array_strategy(
- abi_type: ABIType, registry: StrategyRegistry
- ) -> st.SearchStrategy:
- item_type = abi_type.item_type
- item_type_str = item_type.to_type_str()
- item_strategy = registry.get_strategy(item_type_str)
- last_dim = abi_type.arrlist[-1]
- if len(last_dim) == 0:
- # Is dynamic list. Don't restrict length.
- return st.lists(item_strategy)
- else:
- # Is static list. Restrict length.
- dim_size = last_dim[0]
- return st.lists(item_strategy, min_size=dim_size, max_size=dim_size)
- def get_tuple_strategy(
- abi_type: TupleType, registry: StrategyRegistry
- ) -> st.SearchStrategy:
- component_strategies = [
- registry.get_strategy(comp_abi_type.to_type_str())
- for comp_abi_type in abi_type.components
- ]
- return st.tuples(*component_strategies)
- strategy_registry = StrategyRegistry()
- strategy_registry.register_strategy(
- BaseEquals("uint"),
- get_uint_strategy,
- label="uint",
- )
- strategy_registry.register_strategy(
- BaseEquals("int"),
- get_int_strategy,
- label="int",
- )
- strategy_registry.register_strategy(
- BaseEquals("address", with_sub=False),
- address_strategy,
- label="address",
- )
- strategy_registry.register_strategy(
- BaseEquals("bool", with_sub=False),
- bool_strategy,
- label="bool",
- )
- strategy_registry.register_strategy(
- BaseEquals("ufixed"),
- get_ufixed_strategy,
- label="ufixed",
- )
- strategy_registry.register_strategy(
- BaseEquals("fixed"),
- get_fixed_strategy,
- label="fixed",
- )
- strategy_registry.register_strategy(
- BaseEquals("bytes", with_sub=True),
- get_bytes_strategy,
- label="bytes<M>",
- )
- strategy_registry.register_strategy(
- BaseEquals("bytes", with_sub=False),
- bytes_strategy,
- label="bytes",
- )
- strategy_registry.register_strategy(
- BaseEquals("function", with_sub=False),
- get_bytes_strategy,
- label="function",
- )
- strategy_registry.register_strategy(
- BaseEquals("string", with_sub=False),
- string_strategy,
- label="string",
- )
- strategy_registry.register_strategy(
- has_arrlist,
- get_array_strategy,
- label="has_arrlist",
- )
- strategy_registry.register_strategy(
- is_base_tuple,
- get_tuple_strategy,
- label="is_base_tuple",
- )
- get_abi_strategy = strategy_registry.get_strategy
|