revert back BundleReferenceStrategy

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.",

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)