From 141d74cfb79b3c613a9a6ece25614780211c4596 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Sun, 18 Aug 2024 03:06:35 -0700
Subject: [PATCH] final few tests fix

---
 hypothesis-python/src/hypothesis/stateful.py | 42 ++++++--------------
 1 file changed, 12 insertions(+), 30 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index 0339b6a89d..cb0bc7c4e1 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -461,11 +461,13 @@ def __attrs_post_init__(self):
         for k, v in sorted(self.arguments.items()):
             assert not isinstance(v, BundleReferenceStrategy)
 
-            # if isinstance(v, Bundle):
-            #     bundles.append(v)
-            #     consume = isinstance(v, BundleConsumer)
-            #     # TODO: fix order of initial values reverses when using new Bundle class
-            #     v = BundleReferenceStrategy(v.name, consume=consume)
+            # bundles.append(v)
+
+            if isinstance(v, Bundle):
+                bundles.append(v)
+                consume = isinstance(v, BundleConsumer)
+                # TODO: fix order of initial values reverses when using new Bundle class
+                v = BundleReferenceStrategy(v.name, consume=consume)
 
             self.arguments_strategies[k] = v
         # print("TESTING: self.arguments_strategies: ", self.arguments_strategies)
@@ -520,44 +522,24 @@ def do_draw(self, data):
 class Bundle(SampledFromStrategy):
     def __init__(self, name: str, *, consume: bool = False) -> None:
         self.name = name
-        # self.__reference_strategy = BundleReferenceStrategy(name, consume=consume)
         self.consume = consume
-        # super().__init__()
-        super().__init__(
-            # ("a", "b")
-            (1, 2, 3)
-        )  # Some random items that'll get replaced in do_draw
+        super().__init__((1, 2, 3))  # Some random items that'll get replaced in do_draw
 
     def do_draw(self, data):
         machine = data.draw(self_strategy)
-        # print(
-        #     "New_Bundle do_draw machine: ",
-        #     machine,
-        #     type(machine),
-        #     type(data),
-        #     type(self_strategy),
-        #     self_strategy,
-        # )
+
         bundle = machine.bundle(self.name)
         # print("New_Bundle do_draw bundle: ", bundle)
         if not bundle:
             data.mark_invalid(f"Cannot draw from empty bundle {self.name!r}")
         self.elements = bundle
-        #     if self.elements is None:
-        #     self.initialize_with_bundle(bundle)
-        # self.elements = bundle
 
-        out = super().do_draw(data)
+        reference = super().do_draw(data)
+        out = machine.names_to_values[reference.name]
         if self.consume:
-            bundle.remove(out)  # pragma: no cover # coverage is flaky here
+            bundle.remove(reference)  # pragma: no cover # coverage is flaky here
         return out
 
-        # 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]
-
     def __repr__(self):
         consume = self.consume  # __reference_strategy.consume
         if consume is False: