diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py index 7c60d2752f..f1dfa5205c 100644 --- a/hypothesis-python/src/hypothesis/stateful.py +++ b/hypothesis-python/src/hypothesis/stateful.py @@ -455,8 +455,11 @@ def __attrs_post_init__(self): self.arguments_strategies = {} bundles = [] for k, v in sorted(self.arguments.items()): + assert not isinstance(v, BundleReferenceStrategy) if isinstance(v, Bundle): bundles.append(v) + consume = isinstance(v, BundleConsumer) + v = BundleReferenceStrategy(v.name, consume=consume) self.arguments_strategies[k] = v self.bundles = tuple(bundles) @@ -469,6 +472,26 @@ def __repr__(self) -> str: self_strategy = st.runner() +class BundleReferenceStrategy(SearchStrategy): + def __init__(self, name: str, *, consume: bool = False): + self.name = name + self.consume = consume + + def do_draw(self, data): + machine = data.draw(self_strategy) + bundle = machine.bundle(self.name) + if not bundle: + data.mark_invalid(f"Cannot draw from empty bundle {self.name!r}") + # Shrink towards the right rather than the left. This makes it easier + # to delete data generated earlier, as when the error is towards the + # end there can be a lot of hard to remove padding. + position = data.draw_integer(0, len(bundle) - 1, shrink_towards=len(bundle)) + if self.consume: + return bundle.pop(position) # pragma: no cover # coverage is flaky here + else: + return bundle[position] + + class Bundle(SearchStrategy[Ex]): """A collection of values for use in stateful testing. @@ -495,32 +518,16 @@ def __init__( self, name: str, *, consume: bool = False, draw_references: bool = True ) -> None: self.name = name - self.consume = consume + self.__reference_strategy = BundleReferenceStrategy(name, consume=consume) self.draw_references = draw_references def do_draw(self, data): machine = data.draw(self_strategy) - - bundle = machine.bundle(self.name) - if not bundle: - data.mark_invalid(f"Cannot draw from empty bundle {self.name!r}") - # Shrink towards the right rather than the left. This makes it easier - # to delete data generated earlier, as when the error is towards the - # end there can be a lot of hard to remove padding. - position = data.draw_integer(0, len(bundle) - 1, shrink_towards=len(bundle)) - if self.consume: - reference = bundle.pop( - position - ) # pragma: no cover # coverage is flaky here - else: - reference = bundle[position] - - if self.draw_references: - return reference + reference = data.draw(self.__reference_strategy) return machine.names_to_values[reference.name] def __repr__(self): - consume = self.consume + consume = self.__reference_strategy.consume if consume is False: return f"Bundle(name={self.name!r})" return f"Bundle(name={self.name!r}, {consume=})" @@ -539,11 +546,18 @@ def available(self, data): def flatmap(self, expand): if self.draw_references: return type(self)( - self.name, consume=self.consume, draw_references=False + self.name, + consume=self.__reference_strategy.consume, + draw_references=False, ).flatmap(expand) return super().flatmap(expand) +class BundleConsumer(Bundle[Ex]): + def __init__(self, bundle: Bundle[Ex]) -> None: + super().__init__(bundle.name, consume=True) + + def consumes(bundle: Bundle[Ex]) -> SearchStrategy[Ex]: """When introducing a rule in a RuleBasedStateMachine, this function can be used to mark bundles from which each value used in a step with the @@ -559,10 +573,7 @@ def consumes(bundle: Bundle[Ex]) -> SearchStrategy[Ex]: """ if not isinstance(bundle, Bundle): raise TypeError("Argument to be consumed must be a bundle.") - return type(bundle)( - name=bundle.name, - consume=True, - ) + return BundleConsumer(bundle) @attr.s() @@ -609,7 +620,7 @@ def _convert_targets(targets, target): ) raise InvalidArgument(msg % (t, type(t))) while isinstance(t, Bundle): - if t.consume: + if isinstance(t, BundleConsumer): note_deprecation( f"Using consumes({t.name}) doesn't makes sense in this context. " "This will be an error in a future version of Hypothesis.", diff --git a/hypothesis-python/tests/cover/test_stateful.py b/hypothesis-python/tests/cover/test_stateful.py index 2dd51081cc..ac8473ad13 100644 --- a/hypothesis-python/tests/cover/test_stateful.py +++ b/hypothesis-python/tests/cover/test_stateful.py @@ -24,6 +24,7 @@ reproduce_failure, seed, settings as Settings, + strategies as st, ) from hypothesis.control import current_build_context from hypothesis.database import ExampleDatabase @@ -1339,3 +1340,24 @@ def use_directly(self, bun): Machine.TestCase.settings = Settings(stateful_step_count=5, max_examples=10) run_state_machine_as_test(Machine) + + +def test_use_bundle_within_other_strategies(): + class Class: + def __init__(self, value): + self.value = value + + class Machine(RuleBasedStateMachine): + my_bundle = Bundle("my_bundle") + + @initialize(target=my_bundle) + def set_initial(self, /) -> str: + return "sample text" + + @rule(instance=st.builds(Class, my_bundle)) + def check(self, instance): + assert isinstance(instance, Class) + assert isinstance(instance.value, str) + + Machine.TestCase.settings = Settings(stateful_step_count=5, max_examples=10) + run_state_machine_as_test(Machine)