move more functions to choice.py

Author: tybug

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

@@ -24,7 +24,7 @@
 from hypothesis.errors import ChoiceTooLarge
 from hypothesis.internal.conjecture.floats import float_to_lex, lex_to_float
 from hypothesis.internal.conjecture.utils import identity
-from hypothesis.internal.floats import make_float_clamper, sign_aware_lte
+from hypothesis.internal.floats import float_to_int, make_float_clamper, sign_aware_lte
 from hypothesis.internal.intervalsets import IntervalSet
 
 T = TypeVar("T")
@@ -454,3 +454,42 @@ def choice_permitted(choice: ChoiceT, kwargs: ChoiceKwargsT) -> bool:
         return True
     else:
         raise NotImplementedError(f"unhandled type {type(choice)} with value {choice}")
+
+
+def choices_key(choices: Sequence[ChoiceT]) -> tuple[ChoiceT, ...]:
+    return tuple(choice_key(choice) for choice in choices)
+
+
+def choice_key(choice: ChoiceT) -> ChoiceT:
+    if type(choice) is float:
+        return float_to_int(choice)
+    return choice
+
+
+def choice_equal(choice1: ChoiceT, choice2: ChoiceT) -> bool:
+    assert type(choice1) is type(choice2), (choice1, choice2)
+    return choice_key(choice1) == choice_key(choice2)
+
+
+def choice_kwargs_equal(
+    ir_type: ChoiceNameT, kwargs1: ChoiceKwargsT, kwargs2: ChoiceKwargsT
+) -> bool:
+    return choice_kwargs_key(ir_type, kwargs1) == choice_kwargs_key(ir_type, kwargs2)
+
+
+def choice_kwargs_key(ir_type, kwargs):
+    if ir_type == "float":
+        return (
+            float_to_int(kwargs["min_value"]),
+            float_to_int(kwargs["max_value"]),
+            kwargs["allow_nan"],
+            kwargs["smallest_nonzero_magnitude"],
+        )
+    if ir_type == "integer":
+        return (
+            kwargs["min_value"],
+            kwargs["max_value"],
+            None if kwargs["weights"] is None else tuple(kwargs["weights"]),
+            kwargs["shrink_towards"],
+        )
+    return tuple(kwargs[key] for key in sorted(kwargs))

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

@@ -44,7 +44,11 @@
     FloatKWargs,
     IntegerKWargs,
     StringKWargs,
+    choice_equal,
     choice_from_index,
+    choice_key,
+    choice_kwargs_equal,
+    choice_kwargs_key,
     choice_permitted,
 )
 from hypothesis.internal.conjecture.floats import float_to_lex, lex_to_float
@@ -647,15 +651,15 @@ def trivial(self) -> bool:
 
         if self.ir_type != "float":
             zero_value = choice_from_index(0, self.ir_type, self.kwargs)
-            return ir_value_equal(self.value, zero_value)
+            return choice_equal(self.value, zero_value)
         else:
             kwargs = cast(FloatKWargs, self.kwargs)
             min_value = kwargs["min_value"]
             max_value = kwargs["max_value"]
             shrink_towards = 0.0
 
             if min_value == -math.inf and max_value == math.inf:
-                return ir_value_equal(self.value, shrink_towards)
+                return choice_equal(self.value, shrink_towards)
 
             if (
                 not math.isinf(min_value)
@@ -666,7 +670,7 @@ def trivial(self) -> bool:
                 # one closest to shrink_towards
                 shrink_towards = max(math.ceil(min_value), shrink_towards)
                 shrink_towards = min(math.floor(max_value), shrink_towards)
-                return ir_value_equal(self.value, float(shrink_towards))
+                return choice_equal(self.value, float(shrink_towards))
 
             # the real answer here is "the value in [min_value, max_value] with
             # the lowest denominator when represented as a fraction".
@@ -680,17 +684,17 @@ def __eq__(self, other: object) -> bool:
 
         return (
             self.ir_type == other.ir_type
-            and ir_value_equal(self.value, other.value)
-            and ir_kwargs_equal(self.ir_type, self.kwargs, other.kwargs)
+            and choice_equal(self.value, other.value)
+            and choice_kwargs_equal(self.ir_type, self.kwargs, other.kwargs)
             and self.was_forced == other.was_forced
         )
 
     def __hash__(self) -> int:
         return hash(
             (
                 self.ir_type,
-                ir_value_key(self.value),
-                ir_kwargs_key(self.ir_type, self.kwargs),
+                choice_key(self.value),
+                choice_kwargs_key(self.ir_type, self.kwargs),
                 self.was_forced,
             )
         )
@@ -726,39 +730,6 @@ def ir_size(ir: Iterable[Union[IRNode, NodeTemplate, ChoiceT]]) -> int:
     return size
 
 
-def ir_value_key(v):
-    if type(v) is float:
-        return float_to_int(v)
-    return v
-
-
-def ir_kwargs_key(ir_type, kwargs):
-    if ir_type == "float":
-        return (
-            float_to_int(kwargs["min_value"]),
-            float_to_int(kwargs["max_value"]),
-            kwargs["allow_nan"],
-            kwargs["smallest_nonzero_magnitude"],
-        )
-    if ir_type == "integer":
-        return (
-            kwargs["min_value"],
-            kwargs["max_value"],
-            None if kwargs["weights"] is None else tuple(kwargs["weights"]),
-            kwargs["shrink_towards"],
-        )
-    return tuple(kwargs[key] for key in sorted(kwargs))
-
-
-def ir_value_equal(v1, v2):
-    assert type(v1) is type(v2), (v1, v2)
-    return ir_value_key(v1) == ir_value_key(v2)
-
-
-def ir_kwargs_equal(ir_type, kwargs1, kwargs2):
-    return ir_kwargs_key(ir_type, kwargs1) == ir_kwargs_key(ir_type, kwargs2)
-
-
 @dataclass_transform()
 @attr.s(slots=True)
 class ConjectureResult:
@@ -1969,7 +1940,7 @@ def _pooled_kwargs(self, ir_type, kwargs):
         if self.provider.avoid_realization:
             return kwargs
 
-        key = (ir_type, *ir_kwargs_key(ir_type, kwargs))
+        key = (ir_type, *choice_kwargs_key(ir_type, kwargs))
         try:
             return POOLED_KWARGS_CACHE[key]
         except KeyError:

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

@@ -44,7 +44,7 @@
 )
 from hypothesis.internal.cache import LRUReusedCache
 from hypothesis.internal.compat import NotRequired, TypeAlias, TypedDict, ceil, override
-from hypothesis.internal.conjecture.choice import ChoiceKwargsT, ChoiceT
+from hypothesis.internal.conjecture.choice import ChoiceKwargsT, ChoiceT, choices_key
 from hypothesis.internal.conjecture.data import (
     AVAILABLE_PROVIDERS,
     ConjectureData,
@@ -58,7 +58,6 @@
     Status,
     _Overrun,
     ir_size,
-    ir_value_key,
 )
 from hypothesis.internal.conjecture.datatree import (
     DataTree,
@@ -346,8 +345,8 @@ def __stoppable_test_function(self, data: ConjectureData) -> None:
                     # correct engine.
                     raise
 
-    def _cache_key(self, choices: Sequence[ChoiceT]) -> tuple[int, ...]:
-        return tuple(ir_value_key(choice) for choice in choices)
+    def _cache_key(self, choices: Sequence[ChoiceT]) -> tuple[ChoiceT, ...]:
+        return choices_key(choices)
 
     def _cache(self, data: ConjectureData) -> None:
         result = data.as_result()

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

@@ -16,7 +16,9 @@
 import attr
 
 from hypothesis.internal.conjecture.choice import (
+    choice_equal,
     choice_from_index,
+    choice_key,
     choice_permitted,
     choice_to_index,
 )
@@ -27,8 +29,6 @@
     Status,
     draw_choice,
     ir_size,
-    ir_value_equal,
-    ir_value_key,
 )
 from hypothesis.internal.conjecture.junkdrawer import (
     endswith,
@@ -1082,7 +1082,7 @@ def __changed_nodes(self):
             assert len(prev_nodes) == len(new_nodes)
             for i, (n1, n2) in enumerate(zip(prev_nodes, new_nodes)):
                 assert n1.ir_type == n2.ir_type
-                if not ir_value_equal(n1.value, n2.value):
+                if not choice_equal(n1.value, n2.value):
                     self.__all_changed_nodes.add(i)
 
         return self.__all_changed_nodes
@@ -1295,7 +1295,7 @@ def duplicated_nodes(self):
         """Returns a list of nodes grouped (ir_type, value)."""
         duplicates = defaultdict(list)
         for node in self.nodes:
-            duplicates[(node.ir_type, ir_value_key(node.value))].append(node)
+            duplicates[(node.ir_type, choice_key(node.value))].append(node)
         return list(duplicates.values())
 
     @defines_shrink_pass()
@@ -1437,7 +1437,7 @@ def minimize_nodes(self, nodes):
         # same operation on both basically just works.
         kwargs = nodes[0].kwargs
         assert all(
-            node.ir_type == ir_type and ir_value_equal(node.value, value)
+            node.ir_type == ir_type and choice_equal(node.value, value)
             for node in nodes
         )
 

hypothesis-python/tests/conjecture/test_forced.py

@@ -15,7 +15,8 @@
 import hypothesis.strategies as st
 from hypothesis import HealthCheck, assume, example, given, settings
 from hypothesis.internal.conjecture import utils as cu
-from hypothesis.internal.conjecture.data import ConjectureData, ir_value_equal
+from hypothesis.internal.conjecture.choice import choice_equal
+from hypothesis.internal.conjecture.data import ConjectureData
 from hypothesis.internal.floats import SIGNALING_NAN, SMALLEST_SUBNORMAL
 
 from tests.conjecture.common import fresh_data, ir_types_and_kwargs
@@ -142,13 +143,13 @@ def test_forced_values(ir_type_and_kwargs):
 
     forced = kwargs["forced"]
     data = fresh_data()
-    assert ir_value_equal(getattr(data, f"draw_{ir_type}")(**kwargs), forced)
+    assert choice_equal(getattr(data, f"draw_{ir_type}")(**kwargs), forced)
 
     # now make sure the written buffer reproduces the forced value, even without
     # specifying forced=.
     del kwargs["forced"]
     data = ConjectureData.for_choices(data.choices)
-    assert ir_value_equal(getattr(data, f"draw_{ir_type}")(**kwargs), forced)
+    assert choice_equal(getattr(data, f"draw_{ir_type}")(**kwargs), forced)
 
 
 @pytest.mark.parametrize("sign", [1, -1])

hypothesis-python/tests/conjecture/test_ir.py

@@ -25,6 +25,7 @@
 )
 from hypothesis.errors import StopTest
 from hypothesis.internal.conjecture.choice import (
+    choice_equal,
     choice_from_index,
     choice_permitted,
     choice_to_index,
@@ -36,7 +37,6 @@
     NodeTemplate,
     Status,
     ir_size,
-    ir_value_equal,
 )
 from hypothesis.internal.conjecture.datatree import (
     MAX_CHILDREN_EFFECTIVELY_INFINITE,
@@ -286,7 +286,7 @@ def test_copy_ir_node(node):
     new_value = draw_value(node.ir_type, node.kwargs)
     # if we drew the same value as before, the node should still be equal
     assert (node.copy(with_value=new_value) == node) is (
-        ir_value_equal(new_value, node.value)
+        choice_equal(new_value, node.value)
     )
 
 
@@ -325,9 +325,9 @@ def test_data_with_changed_forced_value(node):
     draw_func = getattr(data, f"draw_{node.ir_type}")
     kwargs = deepcopy(node.kwargs)
     kwargs["forced"] = draw_value(node.ir_type, node.kwargs)
-    assume(not ir_value_equal(kwargs["forced"], node.value))
+    assume(not choice_equal(kwargs["forced"], node.value))
 
-    assert ir_value_equal(draw_func(**kwargs), kwargs["forced"])
+    assert choice_equal(draw_func(**kwargs), kwargs["forced"])
 
 
 # ensure we hit bare-minimum coverage for all ir types.
@@ -374,7 +374,7 @@ def test_data_with_same_forced_value_is_valid(node):
 
     kwargs = deepcopy(node.kwargs)
     kwargs["forced"] = node.value
-    assert ir_value_equal(draw_func(**kwargs), kwargs["forced"])
+    assert choice_equal(draw_func(**kwargs), kwargs["forced"])
 
 
 @given(ir_types_and_kwargs())
@@ -565,7 +565,7 @@ def values(draw):
         return getattr(data, f"draw_{node.ir_type}")(**node.kwargs)
 
     # if we're trivial, then shrinking should produce the same value.
-    assert ir_value_equal(minimal(values()), node.value)
+    assert choice_equal(minimal(values()), node.value)
 
 
 @pytest.mark.parametrize(
@@ -623,7 +623,7 @@ def values(draw):
         return getattr(data, f"draw_{node.ir_type}")(**node.kwargs)
 
     # if we're nontrivial, then shrinking should produce something different.
-    assert not ir_value_equal(minimal(values()), node.value)
+    assert not choice_equal(minimal(values()), node.value)
 
 
 @pytest.mark.parametrize(
@@ -671,7 +671,7 @@ def values(draw):
         data = draw(st.data()).conjecture_data
         return getattr(data, f"draw_{node.ir_type}")(**node.kwargs)
 
-    assert ir_value_equal(minimal(values()), node.value)
+    assert choice_equal(minimal(values()), node.value)
 
 
 @given(ir_nodes())
@@ -785,7 +785,7 @@ def test_choice_index_and_value_are_inverses(ir_type_and_kwargs):
     v = draw_value(ir_type, kwargs)
     index = choice_to_index(v, kwargs)
     note({"v": v, "index": index})
-    ir_value_equal(choice_from_index(index, ir_type, kwargs), v)
+    choice_equal(choice_from_index(index, ir_type, kwargs), v)
 
 
 @pytest.mark.parametrize(
@@ -814,7 +814,7 @@ def test_choice_index_and_value_are_inverses(ir_type_and_kwargs):
 def test_choice_index_and_value_are_inverses_explicit(ir_type, kwargs, choices):
     for choice in choices:
         index = choice_to_index(choice, kwargs)
-        assert ir_value_equal(choice_from_index(index, ir_type, kwargs), choice)
+        assert choice_equal(choice_from_index(index, ir_type, kwargs), choice)
 
 
 @pytest.mark.parametrize(
@@ -852,7 +852,7 @@ def test_drawing_directly_matches_for_choices(nodes):
     data = ConjectureData.for_choices([n.value for n in nodes])
     for node in nodes:
         value = getattr(data, f"draw_{node.ir_type}")(**node.kwargs)
-        assert ir_value_equal(node.value, value)
+        assert choice_equal(node.value, value)
 
 
 def test_draw_directly_explicit():

hypothesis-python/tests/cover/test_database_backend.py

@@ -35,7 +35,7 @@
 )
 from hypothesis.errors import HypothesisWarning
 from hypothesis.internal.compat import WINDOWS
-from hypothesis.internal.conjecture.data import ir_value_equal
+from hypothesis.internal.conjecture.choice import choice_equal
 from hypothesis.stateful import Bundle, RuleBasedStateMachine, rule
 from hypothesis.strategies import binary, lists, tuples
 from hypothesis.utils.conventions import not_set
@@ -485,7 +485,7 @@ def test_ir_nodes_rountrips(nodes1):
     assert len(nodes1) == len(ir2)
 
     for n1, v2 in zip(nodes1, ir2):
-        assert ir_value_equal(n1.value, v2)
+        assert choice_equal(n1.value, v2)
 
     s2 = ir_to_bytes(ir2)
     assert s1 == s2

hypothesis-python/tests/cover/test_float_utils.py

@@ -14,7 +14,7 @@
 import pytest
 
 from hypothesis import example, given, strategies as st
-from hypothesis.internal.conjecture.data import ir_value_equal
+from hypothesis.internal.conjecture.choice import choice_equal
 from hypothesis.internal.floats import (
     count_between_floats,
     float_permitted,
@@ -103,4 +103,4 @@ def test_float_clamper(kwargs, input_value):
         allow_nan=allow_nan,
         smallest_nonzero_magnitude=smallest_nonzero_magnitude,
     ):
-        assert ir_value_equal(input_value, clamped)
+        assert choice_equal(input_value, clamped)