add variable-size width support to the bytes ir node

Author: tybug

hypothesis-python/src/hypothesis/internal/conjecture/data.py

@@ -115,7 +115,8 @@ class StringKWargs(TypedDict):
 
 
 class BytesKWargs(TypedDict):
-    size: int
+    min_size: int
+    max_size: Optional[int]
 
 
 class BooleanKWargs(TypedDict):
@@ -206,7 +207,7 @@ def structural_coverage(label: int) -> StructuralCoverageTag:
 FLOAT_INIT_LOGIC_CACHE = LRUCache(4096)
 POOLED_KWARGS_CACHE = LRUCache(4096)
 
-DRAW_STRING_DEFAULT_MAX_SIZE = 10**10  # "arbitrarily large"
+COLLECTION_DEFAULT_MAX_SIZE = 10**10  # "arbitrarily large"
 
 
 class Example:
@@ -1036,7 +1037,7 @@ def trivial(self):
             return self.value == (minimal_char * self.kwargs["min_size"])
         if self.ir_type == "bytes":
             # smallest size and all-zero value.
-            return len(self.value) == self.kwargs["size"] and not any(self.value)
+            return len(self.value) == self.kwargs["min_size"] and not any(self.value)
 
         raise NotImplementedError(f"unhandled ir_type {self.ir_type}")
 
@@ -1095,7 +1096,11 @@ def ir_value_permitted(value, ir_type, kwargs):
             return False
         return all(ord(c) in kwargs["intervals"] for c in value)
     elif ir_type == "bytes":
-        return len(value) == kwargs["size"]
+        if len(value) < kwargs["min_size"]:
+            return False
+        if kwargs["max_size"] is not None and len(value) > kwargs["max_size"]:
+            return False
+        return True
     elif ir_type == "boolean":
         if kwargs["p"] <= 2 ** (-64):
             return value is False
@@ -1322,7 +1327,12 @@ def draw_string(
 
     @abc.abstractmethod
     def draw_bytes(
-        self, size: int, *, forced: Optional[bytes] = None, fake_forced: bool = False
+        self,
+        min_size: int,
+        max_size: Optional[int],
+        *,
+        forced: Optional[bytes] = None,
+        fake_forced: bool = False,
     ) -> bytes:
         raise NotImplementedError
 
@@ -1611,7 +1621,7 @@ def draw_string(
         fake_forced: bool = False,
     ) -> str:
         if max_size is None:
-            max_size = DRAW_STRING_DEFAULT_MAX_SIZE
+            max_size = COLLECTION_DEFAULT_MAX_SIZE
 
         assert forced is None or min_size <= len(forced) <= max_size
         assert self._cd is not None
@@ -1663,17 +1673,44 @@ def draw_string(
         return "".join(chars)
 
     def draw_bytes(
-        self, size: int, *, forced: Optional[bytes] = None, fake_forced: bool = False
+        self,
+        min_size: int,
+        max_size: Optional[int],
+        *,
+        forced: Optional[bytes] = None,
+        fake_forced: bool = False,
     ) -> bytes:
-        forced_i = None
-        if forced is not None:
-            forced_i = int_from_bytes(forced)
-            size = len(forced)
+        if max_size is None:
+            max_size = COLLECTION_DEFAULT_MAX_SIZE
 
-        assert self._cd is not None
-        return self._cd.draw_bits(
-            8 * size, forced=forced_i, fake_forced=fake_forced
-        ).to_bytes(size, "big")
+        assert forced is None or min_size <= len(forced) <= max_size
+
+        buf = bytearray()
+        average_size = min(
+            max(min_size * 2, min_size + 5),
+            0.5 * (min_size + max_size),
+        )
+        elements = many(
+            self._cd,
+            min_size=min_size,
+            max_size=max_size,
+            average_size=average_size,
+            forced=None if forced is None else len(forced),
+            fake_forced=fake_forced,
+            observe=False,
+        )
+        while elements.more():
+            forced_i: Optional[int] = None
+            if forced is not None:
+                # implicit conversion from bytes to int by indexing here
+                forced_i = forced[elements.count - 1]
+
+            assert self._cd is not None
+            buf += self._cd.draw_bits(
+                8, forced=forced_i, fake_forced=fake_forced
+            ).to_bytes(1, "big")
+
+        return bytes(buf)
 
     def _draw_float(
         self,
@@ -2222,6 +2259,7 @@ def draw_string(
         observe: bool = True,
     ) -> str:
         assert forced is None or min_size <= len(forced)
+        assert min_size >= 0
 
         kwargs: StringKWargs = self._pooled_kwargs(
             "string",
@@ -2255,17 +2293,19 @@ def draw_string(
 
     def draw_bytes(
         self,
-        # TODO move to min_size and max_size here.
-        size: int,
+        min_size: int,
+        max_size: Optional[int],
         *,
         forced: Optional[bytes] = None,
         fake_forced: bool = False,
         observe: bool = True,
     ) -> bytes:
-        assert forced is None or len(forced) == size
-        assert size >= 0
+        assert forced is None or min_size <= len(forced)
+        assert min_size >= 0
 
-        kwargs: BytesKWargs = self._pooled_kwargs("bytes", {"size": size})
+        kwargs: BytesKWargs = self._pooled_kwargs(
+            "bytes", {"min_size": min_size, "max_size": max_size}
+        )
 
         if self.ir_tree_nodes is not None and observe:
             node_value = self._pop_ir_tree_node("bytes", kwargs, forced=forced)

hypothesis-python/src/hypothesis/internal/conjecture/datatree.py

@@ -147,7 +147,7 @@ def _repr_pretty_(self, p, cycle):
 
 
 def compute_max_children(ir_type, kwargs):
-    from hypothesis.internal.conjecture.data import DRAW_STRING_DEFAULT_MAX_SIZE
+    from hypothesis.internal.conjecture.data import COLLECTION_DEFAULT_MAX_SIZE
 
     if ir_type == "integer":
         min_value = kwargs["min_value"]
@@ -178,14 +178,26 @@ def compute_max_children(ir_type, kwargs):
             return 1
         return 2
     elif ir_type == "bytes":
-        return 2 ** (8 * kwargs["size"])
+        min_size = kwargs["min_size"]
+        max_size = kwargs["max_size"]
+
+        if max_size is None:
+            max_size = COLLECTION_DEFAULT_MAX_SIZE
+
+        definitely_too_large = max_size * math.log(2**8) > math.log(
+            MAX_CHILDREN_EFFECTIVELY_INFINITE
+        )
+        if definitely_too_large:
+            return MAX_CHILDREN_EFFECTIVELY_INFINITE
+
+        return sum(2 ** (8 * k) for k in range(min_size, max_size + 1))
     elif ir_type == "string":
         min_size = kwargs["min_size"]
         max_size = kwargs["max_size"]
         intervals = kwargs["intervals"]
 
         if max_size is None:
-            max_size = DRAW_STRING_DEFAULT_MAX_SIZE
+            max_size = COLLECTION_DEFAULT_MAX_SIZE
 
         if len(intervals) == 0:
             # Special-case the empty alphabet to avoid an error in math.log(0).
@@ -306,8 +318,13 @@ def all_children(ir_type, kwargs):
         else:
             yield from [False, True]
     if ir_type == "bytes":
-        size = kwargs["size"]
-        yield from (int_to_bytes(i, size) for i in range(2 ** (8 * size)))
+        min_size = kwargs["min_size"]
+        max_size = kwargs["max_size"]
+
+        size = min_size
+        while size <= max_size:
+            yield from (int_to_bytes(i, size) for i in range(2 ** (8 * size)))
+            size += 1
     if ir_type == "string":
         min_size = kwargs["min_size"]
         max_size = kwargs["max_size"]

hypothesis-python/src/hypothesis/internal/conjecture/shrinker.py

@@ -1075,10 +1075,9 @@ def try_shrinking_nodes(self, nodes, n):
                 return False  # pragma: no cover
 
             if node.ir_type in {"string", "bytes"}:
-                size_kwarg = "min_size" if node.ir_type == "string" else "size"
                 # if the size *increased*, we would have to guess what to pad with
                 # in order to try fixing up this attempt. Just give up.
-                if node.kwargs[size_kwarg] <= attempt_kwargs[size_kwarg]:
+                if node.kwargs["min_size"] <= attempt_kwargs["min_size"]:
                     return False
                 # the size decreased in our attempt. Try again, but replace with
                 # the min_size that we would have gotten, and truncate the value
@@ -1089,7 +1088,7 @@ def try_shrinking_nodes(self, nodes, n):
                         initial_attempt[node.index].copy(
                             with_kwargs=attempt_kwargs,
                             with_value=initial_attempt[node.index].value[
-                                : attempt_kwargs[size_kwarg]
+                                : attempt_kwargs["min_size"]
                             ],
                         )
                     ]

hypothesis-python/src/hypothesis/internal/conjecture/shrinking/bytes.py

@@ -8,17 +8,19 @@
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at https://mozilla.org/MPL/2.0/.
 
-from hypothesis.internal.compat import int_from_bytes, int_to_bytes
+from hypothesis.internal.conjecture.shrinking.collection import Collection
 from hypothesis.internal.conjecture.shrinking.integer import Integer
 
 
-class Bytes(Integer):
+class Bytes(Collection):
     def __init__(self, initial, predicate, **kwargs):
         # shrink by interpreting the bytes as an integer.
         # move to Collection.shrink when we support variable-size bytes,
         # because b'\x00\x02' could shrink to either b'\x00\x01' or b'\x02'.
         super().__init__(
-            int_from_bytes(initial),
-            lambda n: predicate(int_to_bytes(n, len(initial))),
+            # implicit conversion from bytes to list of integers here
+            list(initial),
+            lambda val: predicate(bytes(val)),
+            ElementShrinker=Integer,
             **kwargs,
         )

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

@@ -12,6 +12,7 @@
 import re
 import warnings
 from functools import lru_cache, partial
+from typing import Optional
 
 from hypothesis.errors import HypothesisWarning, InvalidArgument
 from hypothesis.internal import charmap
@@ -20,7 +21,6 @@
 from hypothesis.internal.reflection import get_pretty_function_description
 from hypothesis.strategies._internal.collections import ListStrategy
 from hypothesis.strategies._internal.lazy import unwrap_strategies
-from hypothesis.strategies._internal.numbers import IntegersStrategy
 from hypothesis.strategies._internal.strategies import (
     OneOfStrategy,
     SampledFromStrategy,
@@ -224,9 +224,13 @@ def _string_filter_rewrite(self, kind, condition):
             stacklevel=2,
         )
 
-    elems = unwrap_strategies(self.element_strategy)
     if (
-        (kind is bytes or isinstance(elems, OneCharStringStrategy))
+        (
+            kind is bytes
+            or isinstance(
+                unwrap_strategies(self.element_strategy), OneCharStringStrategy
+            )
+        )
         and isinstance(pattern := getattr(condition, "__self__", None), re.Pattern)
         and isinstance(pattern.pattern, kind)
     ):
@@ -331,15 +335,13 @@ def _identifier_characters():
     return id_start, id_continue
 
 
-class BytesStrategy(ListStrategy):
-    def __init__(self, min_size, max_size):
-        super().__init__(IntegersStrategy(0, 255), min_size=min_size, max_size=max_size)
+class BytesStrategy(SearchStrategy):
+    def __init__(self, min_size: int, max_size: Optional[int]):
+        self.min_size = min_size
+        self.max_size = max_size if max_size is not None else float("inf")
 
     def do_draw(self, data):
-        # TODO: refactor the underlying provider to support variable-length bytes
-        if self.min_size == self.max_size:
-            return bytes(data.draw_bytes(self.min_size))
-        return bytes(super().do_draw(data))
+        return data.draw_bytes(self.min_size, self.max_size)
 
     _nonempty_filters = (
         *ListStrategy._nonempty_filters,
@@ -353,4 +355,4 @@ def do_draw(self, data):
     def filter(self, condition):
         if (new := _string_filter_rewrite(self, bytes, condition)) is not None:
             return new
-        return super().filter(condition)
+        return ListStrategy.filter(self, condition)

hypothesis-python/tests/conjecture/common.py

@@ -170,49 +170,56 @@ def draw_integer_kwargs(
 
 
 @st.composite
-def draw_string_kwargs(draw, *, use_min_size=True, use_max_size=True, use_forced=False):
-    # TODO also sample empty intervals, ie remove this min_size, once we handle empty
-    # pseudo-choices in the ir
-    interval_set = draw(intervals(min_size=1))
-    forced = (
-        draw(TextStrategy(OneCharStringStrategy(interval_set))) if use_forced else None
-    )
-
+def _collection_kwargs(draw, *, forced, use_min_size=True, use_max_size=True):
     min_size = 0
     max_size = None
+    # collections are quite expensive in entropy. cap to avoid overruns.
+    cap = 50
 
     if use_min_size:
-        # cap to some reasonable min size to avoid overruns.
-        n = 100
-        if forced is not None:
-            n = min(n, len(forced))
-
-        min_size = draw(st.integers(0, n))
+        min_size = draw(
+            st.integers(0, min(len(forced), cap) if forced is not None else cap)
+        )
 
     if use_max_size:
-        n = min_size if forced is None else max(min_size, len(forced))
-        max_size = draw(st.integers(min_value=n))
+        max_size = draw(
+            st.integers(
+                min_value=min_size if forced is None else max(min_size, len(forced))
+            )
+        )
         # cap to some reasonable max size to avoid overruns.
         max_size = min(max_size, min_size + 100)
 
-    return {
-        "intervals": interval_set,
-        "min_size": min_size,
-        "max_size": max_size,
-        "forced": forced,
-    }
+    return {"min_size": min_size, "max_size": max_size}
 
 
 @st.composite
-def draw_bytes_kwargs(draw, *, use_forced=False):
-    forced = draw(st.binary()) if use_forced else None
-    # be reasonable with the number of bytes we ask for. We only have BUFFER_SIZE
-    # to work with before we overrun.
-    size = (
-        draw(st.integers(min_value=0, max_value=100)) if forced is None else len(forced)
+def draw_string_kwargs(draw, *, use_min_size=True, use_max_size=True, use_forced=False):
+    # TODO also sample empty intervals, ie remove this min_size, once we handle empty
+    # pseudo-choices in the ir
+    interval_set = draw(intervals(min_size=1))
+    forced = (
+        draw(TextStrategy(OneCharStringStrategy(interval_set))) if use_forced else None
+    )
+    kwargs = draw(
+        _collection_kwargs(
+            forced=forced, use_min_size=use_min_size, use_max_size=use_max_size
+        )
     )
 
-    return {"size": size, "forced": forced}
+    return {"intervals": interval_set, "forced": forced, **kwargs}
+
+
+@st.composite
+def draw_bytes_kwargs(draw, *, use_min_size=True, use_max_size=True, use_forced=False):
+    forced = draw(st.binary()) if use_forced else None
+
+    kwargs = draw(
+        _collection_kwargs(
+            forced=forced, use_min_size=use_min_size, use_max_size=use_max_size
+        )
+    )
+    return {"forced": forced, **kwargs}
 
 
 @st.composite