merge fixed_dictionaries strategies

Author: tybug

hypothesis-python/RELEASE.rst

@@ -0,0 +1,3 @@
+RELEASE_TYPE: patch
+
+This patch improves the string representation of :func:`~hypothesis.strategies.fixed_dictionaries`.

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

@@ -10,8 +10,9 @@
 
 import copy
 from collections.abc import Iterable
-from typing import Any, overload
+from typing import Any, Optional, overload
 
+from hypothesis import strategies as st
 from hypothesis.errors import InvalidArgument
 from hypothesis.internal.conjecture import utils as cu
 from hypothesis.internal.conjecture.engine import BUFFER_SIZE
@@ -24,7 +25,6 @@
     T4,
     T5,
     Ex,
-    MappedStrategy,
     SearchStrategy,
     T,
     check_strategy,
@@ -306,50 +306,32 @@ def do_draw(self, data):
         return result
 
 
-class FixedKeysDictStrategy(MappedStrategy):
+class FixedDictStrategy(SearchStrategy):
     """A strategy which produces dicts with a fixed set of keys, given a
     strategy for each of their equivalent values.
 
     e.g. {'foo' : some_int_strategy} would generate dicts with the single
     key 'foo' mapping to some integer.
     """
 
-    def __init__(self, strategy_dict):
-        dict_type = type(strategy_dict)
-        self.keys = tuple(strategy_dict.keys())
-        super().__init__(
-            strategy=TupleStrategy(strategy_dict[k] for k in self.keys),
-            pack=lambda value: dict_type(zip(self.keys, value)),
+    def __init__(
+        self,
+        mapping: dict[T, SearchStrategy[Ex]],
+        *,
+        optional: Optional[dict[T, SearchStrategy[Ex]]],
+    ):
+        dict_type = type(mapping)
+        keys = tuple(mapping.keys())
+        self.fixed = st.tuples(*[mapping[k] for k in keys]).map(
+            lambda value: dict_type(zip(keys, value))
         )
-
-    def calc_is_empty(self, recur):
-        return recur(self.mapped_strategy)
-
-    def __repr__(self):
-        return f"FixedKeysDictStrategy({self.keys!r}, {self.mapped_strategy!r})"
-
-
-class FixedAndOptionalKeysDictStrategy(SearchStrategy):
-    """A strategy which produces dicts with a fixed set of keys, given a
-    strategy for each of their equivalent values.
-
-    e.g. {'foo' : some_int_strategy} would generate dicts with the single
-    key 'foo' mapping to some integer.
-    """
-
-    def __init__(self, strategy_dict, optional):
-        self.required = strategy_dict
-        self.fixed = FixedKeysDictStrategy(strategy_dict)
         self.optional = optional
 
-    def calc_is_empty(self, recur):
-        return recur(self.fixed)
-
-    def __repr__(self):
-        return f"FixedAndOptionalKeysDictStrategy({self.required!r}, {self.optional!r})"
-
     def do_draw(self, data):
-        result = data.draw(self.fixed)
+        value = data.draw(self.fixed)
+        if self.optional is None:
+            return value
+
         remaining = [k for k, v in self.optional.items() if not v.is_empty]
         should_draw = cu.many(
             data, min_size=0, max_size=len(remaining), average_size=len(remaining) / 2
@@ -358,5 +340,13 @@ def do_draw(self, data):
             j = data.draw_integer(0, len(remaining) - 1)
             remaining[-1], remaining[j] = remaining[j], remaining[-1]
             key = remaining.pop()
-            result[key] = data.draw(self.optional[key])
-        return result
+            value[key] = data.draw(self.optional[key])
+        return value
+
+    def calc_is_empty(self, recur):
+        return recur(self.fixed)
+
+    def __repr__(self):
+        if self.optional is not None:
+            return f"fixed_dictionaries({self.keys!r}, optional={self.optional!r})"
+        return f"fixed_dictionaries({self.keys!r})"

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

@@ -105,8 +105,7 @@
 )
 from hypothesis.strategies._internal import SearchStrategy, check_strategy
 from hypothesis.strategies._internal.collections import (
-    FixedAndOptionalKeysDictStrategy,
-    FixedKeysDictStrategy,
+    FixedDictStrategy,
     ListStrategy,
     TupleStrategy,
     UniqueListStrategy,
@@ -510,6 +509,7 @@ def fixed_dictionaries(
     check_type(dict, mapping, "mapping")
     for k, v in mapping.items():
         check_strategy(v, f"mapping[{k!r}]")
+
     if optional is not None:
         check_type(dict, optional, "optional")
         for k, v in optional.items():
@@ -524,8 +524,8 @@ def fixed_dictionaries(
                 "The following keys were in both mapping and optional, "
                 f"which is invalid: {set(mapping) & set(optional)!r}"
             )
-        return FixedAndOptionalKeysDictStrategy(mapping, optional)
-    return FixedKeysDictStrategy(mapping)
+
+    return FixedDictStrategy(mapping, optional=optional)
 
 
 @cacheable