Merge Bundle as subclass of SampledFromStrategy

hypothesis-python/RELEASE.rst

@@ -0,0 +1,7 @@
+RELEASE_TYPE: minor
+
+This release changes :class:`hypothesis.stateful.Bundle` to use the internals of
+:func:`~hypothesis.strategies.sampled_from`, improving the `filter` and `map` methods.
+In addition to performance improvements, you can now ``consumes(some_bundle).filter(...)``!
+
+Thanks to Reagan Lee for this feature (:issue:`3944`).

hypothesis-python/docs/data.rst

@@ -56,6 +56,7 @@ Possibly the most important one to be aware of is
 produced by strategies earlier in its argument list. Most of the others should
 largely "do the right thing" without you having to think about it.
 
+.. _adapting-strategies:
 
 ~~~~~~~~~~~~~~~~~~~
 Adapting strategies

hypothesis-python/docs/stateful.rst

@@ -372,3 +372,37 @@ This is not recommended as it bypasses some important internal functions,
 including reporting of statistics such as runtimes and :func:`~hypothesis.event`
 calls.  It was originally added to support custom ``__init__`` methods, but
 you can now use :func:`~hypothesis.stateful.initialize` rules instead.
+
+-------------------------
+Adapting stateful testing 
+-------------------------
+
+You can adapt stateful tests in a similar way to adapting regular strategies. 
+See :ref:`this section <adapting-strategies>` for more details on how each of these strategies work.
+
+.. code:: python
+
+ class Machine(RuleBasedStateMachine):
+     bun = Bundle("bun")
+
+     @initialize(target=buns)
+     def create_bun(self):
+         return multiple(1, 2)
+
+     @rule(bun=buns.map(lambda x: -x))
+     def use_map(self, bun):
+         print("Bun example is negative.")
+
+     @rule(bun=buns.filter(lambda x: x > 1))
+     def use_filter(self, bun):
+         print("Bun example is greater than 1.")
+
+     @rule(bun=buns.flatmap(lambda n: lists(lists(integers(), min_size=n, max_size=n))))
+     def use_flatmap(self, bun):
+         print("Bun example is a unique structure.")
+
+     @rule(bun=buns.map(lambda x: -x).filter(lambda x: x < -1))
+     def use_chain_strategies(self, bun):
+         print("Bun example uses chained strategies!")
+
+

hypothesis-python/src/hypothesis/stateful.py

@@ -56,8 +56,10 @@
     Ex,
     Ex_Inv,
     OneOfStrategy,
+    SampledFromStrategy,
     SearchStrategy,
     check_strategy,
+    filter_not_satisfied,
 )
 from hypothesis.vendor.pretty import RepresentationPrinter
 
@@ -184,12 +186,12 @@ def output(s):
                 try:
                     data = dict(data)
                     for k, v in list(data.items()):
-                        if isinstance(v, VarReference):
-                            data[k] = machine.names_to_values[v.name]
+                        if isinstance(v, VarReferenceMapping):
+                            data[k] = v.value
                         elif isinstance(v, list) and all(
-                            isinstance(item, VarReference) for item in v
+                            isinstance(item, VarReferenceMapping) for item in v
                         ):
-                            data[k] = [machine.names_to_values[item.name] for item in v]
+                            data[k] = [item.value for item in v]
 
                     label = f"execute:rule:{rule.function.__name__}"
                     start = perf_counter()
@@ -292,12 +294,12 @@ def __init__(self) -> None:
             )
 
     def _pretty_print(self, value):
-        if isinstance(value, VarReference):
-            return value.name
+        if isinstance(value, VarReferenceMapping):
+            return value.reference.name
         elif isinstance(value, list) and all(
-            isinstance(item, VarReference) for item in value
+            isinstance(item, VarReferenceMapping) for item in value
         ):
-            return "[" + ", ".join([item.name for item in value]) + "]"
+            return "[" + ", ".join([item.reference.name for item in value]) + "]"
         self.__stream.seek(0)
         self.__stream.truncate(0)
         self.__printer.output_width = 0
@@ -458,8 +460,12 @@ def __attrs_post_init__(self):
             assert not isinstance(v, BundleReferenceStrategy)
             if isinstance(v, Bundle):
                 bundles.append(v)
-                consume = isinstance(v, BundleConsumer)
-                v = BundleReferenceStrategy(v.name, consume=consume)
+                v = BundleReferenceStrategy(
+                    v.name,
+                    consume=v.consume,
+                    repr_=v.repr_,
+                    transformations=v.transformations,
+                )
             self.arguments_strategies[k] = v
         self.bundles = tuple(bundles)
 
@@ -472,24 +478,61 @@ def __repr__(self) -> str:
 self_strategy = st.runner()
 
 
-class BundleReferenceStrategy(SearchStrategy):
-    def __init__(self, name: str, *, consume: bool = False):
+class BundleReferenceStrategy(SampledFromStrategy[Ex]):
+
+    def __init__(
+        self,
+        name: str,
+        *,
+        consume: bool = False,
+        repr_: Optional[str] = None,
+        transformations: Iterable[tuple[str, Callable]] = (),
+    ):
         self.name = name
         self.consume = consume
+        super().__init__(
+            [...],
+            repr_=repr_,
+            transformations=transformations,
+        )  # Some random items that'll get replaced in do_draw
+
+    def get_transformed_value(self, reference):
+        assert isinstance(reference, VarReference)
+        return self._transform(self.machine.names_to_values.get(reference.name))
+
+    def get_element(self, i):
+        idx = self.elements[i]
+        assert isinstance(idx, int)
+        reference = self.bundle[idx]
+        value = self.get_transformed_value(reference)
+        if value is filter_not_satisfied:
+            return filter_not_satisfied
+        return idx
 
     def do_draw(self, data):
-        machine = data.draw(self_strategy)
-        bundle = machine.bundle(self.name)
-        if not bundle:
+        self.machine = data.draw(self_strategy)
+        self.bundle = self.machine.bundle(self.name)
+        if not self.bundle:
             data.mark_invalid(f"Cannot draw from empty bundle {self.name!r}")
+
+        # We use both self.bundle and self.elements to make sure an index is
+        # used to safely pop.
+
         # 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))
+        self.elements = range(len(self.bundle))[::-1]
+
+        position = super().do_draw(data)
+        reference = self.bundle[position]
         if self.consume:
-            return bundle.pop(position)  # pragma: no cover  # coverage is flaky here
-        else:
-            return bundle[position]
+            self.bundle.pop(position)  # pragma: no cover # coverage is flaky here
+
+        value = self.get_transformed_value(reference)
+
+        # We need both reference and the value itself to pretty-print deterministically
+        # and maintain any transformations that is bundle-specific
+        return VarReferenceMapping(reference, value)
 
 
 class Bundle(SearchStrategy[Ex]):
@@ -515,22 +558,62 @@ class MyStateMachine(RuleBasedStateMachine):
     """
 
     def __init__(
-        self, name: str, *, consume: bool = False, draw_references: bool = True
+        self,
+        name: str,
+        *,
+        consume: bool = False,
+        repr_: Optional[str] = None,
+        transformations: Iterable[tuple[str, Callable]] = (),
     ) -> None:
         self.name = name
-        self.__reference_strategy = BundleReferenceStrategy(name, consume=consume)
-        self.draw_references = draw_references
+        self.__reference_strategy = BundleReferenceStrategy(
+            name, consume=consume, repr_=repr_, transformations=transformations
+        )
+
+    @property
+    def consume(self):
+        return self.__reference_strategy.consume
+
+    @property
+    def repr_(self):
+        return self.__reference_strategy.repr_
+
+    @property
+    def transformations(self):
+        return self.__reference_strategy._transformations
 
     def do_draw(self, data):
-        machine = data.draw(self_strategy)
-        reference = data.draw(self.__reference_strategy)
-        return machine.names_to_values[reference.name]
+        self.machine = data.draw(self_strategy)
+        var_reference = data.draw(self.__reference_strategy)
+        assert type(var_reference) is VarReferenceMapping
+        return var_reference.value
+
+    def filter(self, condition):
+        return type(self)(
+            self.name,
+            consume=self.__reference_strategy.consume,
+            transformations=(
+                *self.__reference_strategy._transformations,
+                ("filter", condition),
+            ),
+            repr_=self.__reference_strategy.repr_,
+        )
+
+    def map(self, pack):
+        return type(self)(
+            self.name,
+            consume=self.__reference_strategy.consume,
+            transformations=(
+                *self.__reference_strategy._transformations,
+                ("map", pack),
+            ),
+            repr_=self.__reference_strategy.repr_,
+        )
 
     def __repr__(self):
-        consume = self.__reference_strategy.consume
-        if consume is False:
+        if self.consume is False:
             return f"Bundle(name={self.name!r})"
-        return f"Bundle(name={self.name!r}, {consume=})"
+        return f"Bundle(name={self.name!r}, {self.consume=})"
 
     def calc_is_empty(self, recur):
         # We assume that a bundle will grow over time
@@ -543,20 +626,6 @@ 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.__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
@@ -573,7 +642,12 @@ def consumes(bundle: Bundle[Ex]) -> SearchStrategy[Ex]:
     """
     if not isinstance(bundle, Bundle):
         raise TypeError("Argument to be consumed must be a bundle.")
-    return BundleConsumer(bundle)
+    return type(bundle)(
+        name=bundle.name,
+        consume=True,
+        transformations=bundle.transformations,
+        repr_=bundle.repr_,
+    )
 
 
 @attr.s()
@@ -620,7 +694,7 @@ def _convert_targets(targets, target):
                 )
             raise InvalidArgument(msg % (t, type(t)))
         while isinstance(t, Bundle):
-            if isinstance(t, BundleConsumer):
+            if t.consume:
                 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.",
@@ -839,6 +913,12 @@ class VarReference:
     name = attr.ib()
 
 
+@attr.s()
+class VarReferenceMapping:
+    reference: VarReference = attr.ib()
+    value: Any = attr.ib()
+
+
 # There are multiple alternatives for annotating the `precond` type, all of them
 # have drawbacks. See https://github.com/HypothesisWorks/hypothesis/pull/3068#issuecomment-906642371
 def precondition(precond: Callable[[Any], bool]) -> Callable[[TestFunc], TestFunc]:

hypothesis-python/src/hypothesis/strategies/_internal/strategies.py

@@ -477,7 +477,7 @@ def is_simple_data(value):
         return False
 
 
-class SampledFromStrategy(SearchStrategy):
+class SampledFromStrategy(SearchStrategy[Ex]):
     """A strategy which samples from a set of elements. This is essentially
     equivalent to using a OneOfStrategy over Just strategies but may be more
     efficient and convenient.
@@ -528,7 +528,7 @@ def calc_is_cacheable(self, recur):
         return is_simple_data(self.elements)
 
     def _transform(self, element):
-        # Used in UniqueSampledListStrategy
+        # Used in UniqueSampledListStrategy and BundleStrategy
         for name, f in self._transformations:
             if name == "map":
                 result = f(element)
@@ -602,7 +602,11 @@ def do_filtered_draw(self, data):
                     if len(allowed) > speculative_index:
                         # Early-exit case: We reached the speculative index, so
                         # we just return the corresponding element.
-                        data.draw_integer(0, len(self.elements) - 1, forced=i)
+                        data.draw_integer(
+                            0,
+                            len(self.elements) - 1,
+                            forced=i,
+                        )
                         return element
 
         # The speculative index didn't work out, but at this point we've built