Fix Bundle's flatmap functionality

hypothesis-python/RELEASE.rst

@@ -0,0 +1,3 @@
+RELEASE_TYPE: patch
+
+This release fixes a regression where :class:`hypothesis.stateful.Bundle` did not work properly with :ref:`flatmap <flatmap>` functionality (:issue:`4128`).

hypothesis-python/src/hypothesis/stateful.py

@@ -501,9 +501,12 @@ class MyStateMachine(RuleBasedStateMachine):
     drawn from this bundle will be consumed (as above) when requested.
     """
 
-    def __init__(self, name: str, *, consume: bool = False) -> None:
+    def __init__(
+        self, name: str, *, consume: bool = False, draw_references: bool = True
+    ) -> None:
         self.name = name
         self.consume = consume
+        self.draw_references = draw_references
 
     def do_draw(self, data):
         machine = data.draw(self_strategy)
@@ -522,7 +525,9 @@ def do_draw(self, data):
         else:
             reference = bundle[position]
 
-        return reference
+        if self.draw_references:
+            return reference
+        return machine.names_to_values[reference.name]
 
     def __repr__(self):
         consume = self.consume
@@ -541,6 +546,13 @@ def available(self, data):
         machine = data.draw(self_strategy)
         return bool(machine.bundle(self.name))
 
+    def flatmap(self, expand):
+        if self.draw_references:
+            return type(self)(
+                self.name, consume=self.consume, draw_references=False
+            ).flatmap(expand)
+        return super().flatmap(expand)
+
 
 def consumes(bundle: Bundle[Ex]) -> SearchStrategy[Ex]:
     """When introducing a rule in a RuleBasedStateMachine, this function can

hypothesis-python/tests/cover/test_stateful.py

@@ -1318,3 +1318,24 @@ def rule1(self, data):
 
 
 TestLotsOfEntropyPerStepMachine = LotsOfEntropyPerStepMachine.TestCase
+
+
+def test_flatmap():
+    class Machine(RuleBasedStateMachine):
+        buns = Bundle("buns")
+
+        @initialize(target=buns)
+        def create_bun(self):
+            return 0
+
+        @rule(target=buns, bun=buns.flatmap(lambda x: just(x + 1)))
+        def use_flatmap(self, bun):
+            assert isinstance(bun, int)
+            return bun
+
+        @rule(bun=buns)
+        def use_directly(self, bun):
+            assert isinstance(bun, int)
+
+    Machine.TestCase.settings = Settings(stateful_step_count=5, max_examples=10)
+    run_state_machine_as_test(Machine)