Move oneway strategies to hypothesis_helpers.py

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -4,7 +4,7 @@
 from operator import mul
 from typing import Any, List, NamedTuple, Optional, Sequence, Tuple, Union
 
-from hypothesis import assume
+from hypothesis import assume, reject
 from hypothesis.strategies import (SearchStrategy, booleans, composite, floats,
                                    integers, just, lists, none, one_of,
                                    sampled_from, shared)
@@ -99,6 +99,46 @@ def mutually_promotable_dtypes(
     return one_of(strats).map(tuple)
 
 
+class OnewayPromotableDtypes(NamedTuple):
+    input_dtype: DataType
+    result_dtype: DataType
+
+
+@composite
+def oneway_promotable_dtypes(
+    draw, dtypes: Sequence[DataType]
+) -> SearchStrategy[OnewayPromotableDtypes]:
+    """Return a strategy for input dtypes that promote to result dtypes."""
+    d1, d2 = draw(mutually_promotable_dtypes(dtypes=dtypes))
+    result_dtype = dh.result_type(d1, d2)
+    if d1 == result_dtype:
+        return OnewayPromotableDtypes(d2, d1)
+    elif d2 == result_dtype:
+        return OnewayPromotableDtypes(d1, d2)
+    else:
+        reject()
+
+
+class OnewayBroadcastableShapes(NamedTuple):
+    input_shape: Shape
+    result_shape: Shape
+
+
+@composite
+def oneway_broadcastable_shapes(draw) -> SearchStrategy[OnewayBroadcastableShapes]:
+    """Return a strategy for input shapes that broadcast to result shapes."""
+    result_shape = draw(shapes(min_side=1))
+    input_shape = draw(
+        xps.broadcastable_shapes(
+            result_shape,
+            # Override defaults so bad shapes are less likely to be generated.
+            max_side=None if result_shape == () else max(result_shape),
+            max_dims=len(result_shape),
+        ).filter(lambda s: sh.broadcast_shapes(result_shape, s) == result_shape)
+    )
+    return OnewayBroadcastableShapes(input_shape, result_shape)
+
+
 # shared() allows us to draw either the function or the function name and they
 # will both correspond to the same function.
 

array_api_tests/meta/test_utils.py

@@ -4,15 +4,12 @@
 
 from .. import _array_module as xp
 from .. import dtype_helpers as dh
+from .. import hypothesis_helpers as hh
 from .. import shape_helpers as sh
 from .. import xps
 from ..test_creation_functions import frange
 from ..test_manipulation_functions import roll_ndindex
-from ..test_operators_and_elementwise_functions import (
-    mock_int_dtype,
-    oneway_broadcastable_shapes,
-    oneway_promotable_dtypes,
-)
+from ..test_operators_and_elementwise_functions import mock_int_dtype
 
 
 @pytest.mark.parametrize(
@@ -115,11 +112,11 @@ def test_int_to_dtype(x, dtype):
     assert mock_int_dtype(x, dtype) == d
 
 
-@given(oneway_promotable_dtypes(dh.all_dtypes))
+@given(hh.oneway_promotable_dtypes(dh.all_dtypes))
 def test_oneway_promotable_dtypes(D):
     assert D.result_dtype == dh.result_type(*D)
 
 
-@given(oneway_broadcastable_shapes())
+@given(hh.oneway_broadcastable_shapes())
 def test_oneway_broadcastable_shapes(S):
     assert S.result_shape == sh.broadcast_shapes(*S)

array_api_tests/test_array_object.py

@@ -12,7 +12,6 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from .test_operators_and_elementwise_functions import oneway_promotable_dtypes
 from .typing import DataType, Index, Param, Scalar, ScalarType, Shape
 
 pytestmark = pytest.mark.ci
@@ -108,7 +107,7 @@ def test_getitem(shape, dtype, data):
 
 @given(
     shape=hh.shapes(),
-    dtypes=oneway_promotable_dtypes(dh.all_dtypes),
+    dtypes=hh.oneway_promotable_dtypes(dh.all_dtypes),
     data=st.data(),
 )
 def test_setitem(shape, dtypes, data):

array_api_tests/test_creation_functions.py

@@ -12,7 +12,6 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from .test_operators_and_elementwise_functions import oneway_promotable_dtypes
 from .typing import DataType, Scalar
 
 pytestmark = pytest.mark.ci
@@ -256,7 +255,7 @@ def scalar_eq(s1: Scalar, s2: Scalar) -> bool:
 
 @given(
     shape=hh.shapes(),
-    dtypes=oneway_promotable_dtypes(dh.all_dtypes),
+    dtypes=hh.oneway_promotable_dtypes(dh.all_dtypes),
     data=st.data(),
 )
 def test_asarray_arrays(shape, dtypes, data):

array_api_tests/test_operators_and_elementwise_functions.py

@@ -8,9 +8,8 @@
 from typing import Callable, List, NamedTuple, Optional, Sequence, TypeVar, Union
 
 import pytest
-from hypothesis import assume, given
+from hypothesis import assume, given, reject
 from hypothesis import strategies as st
-from hypothesis.control import reject
 
 from . import _array_module as xp, api_version
 from . import array_helpers as ah
@@ -41,46 +40,6 @@ def all_floating_dtypes() -> st.SearchStrategy[DataType]:
     return strat
 
 
-class OnewayPromotableDtypes(NamedTuple):
-    input_dtype: DataType
-    result_dtype: DataType
-
-
-@st.composite
-def oneway_promotable_dtypes(
-    draw, dtypes: Sequence[DataType]
-) -> st.SearchStrategy[OnewayPromotableDtypes]:
-    """Return a strategy for input dtypes that promote to result dtypes."""
-    d1, d2 = draw(hh.mutually_promotable_dtypes(dtypes=dtypes))
-    result_dtype = dh.result_type(d1, d2)
-    if d1 == result_dtype:
-        return OnewayPromotableDtypes(d2, d1)
-    elif d2 == result_dtype:
-        return OnewayPromotableDtypes(d1, d2)
-    else:
-        reject()
-
-
-class OnewayBroadcastableShapes(NamedTuple):
-    input_shape: Shape
-    result_shape: Shape
-
-
-@st.composite
-def oneway_broadcastable_shapes(draw) -> st.SearchStrategy[OnewayBroadcastableShapes]:
-    """Return a strategy for input shapes that broadcast to result shapes."""
-    result_shape = draw(hh.shapes(min_side=1))
-    input_shape = draw(
-        xps.broadcastable_shapes(
-            result_shape,
-            # Override defaults so bad shapes are less likely to be generated.
-            max_side=None if result_shape == () else max(result_shape),
-            max_dims=len(result_shape),
-        ).filter(lambda s: sh.broadcast_shapes(result_shape, s) == result_shape)
-    )
-    return OnewayBroadcastableShapes(input_shape, result_shape)
-
-
 def mock_int_dtype(n: int, dtype: DataType) -> int:
     """Returns equivalent of `n` that mocks `dtype` behaviour."""
     nbits = dh.dtype_nbits[dtype]
@@ -557,7 +516,7 @@ def make_binary_params(
 ) -> List[Param[BinaryParamContext]]:
     if hh.FILTER_UNDEFINED_DTYPES:
         dtypes = [d for d in dtypes if not isinstance(d, xp._UndefinedStub)]
-    shared_oneway_dtypes = st.shared(oneway_promotable_dtypes(dtypes))
+    shared_oneway_dtypes = st.shared(hh.oneway_promotable_dtypes(dtypes))
     left_dtypes = shared_oneway_dtypes.map(lambda D: D.result_dtype)
     right_dtypes = shared_oneway_dtypes.map(lambda D: D.input_dtype)
 
@@ -576,7 +535,7 @@ def make_param(
             right_strat = right_dtypes.flatmap(lambda d: xps.from_dtype(d, **finite_kw))
         else:
             if func_type is FuncType.IOP:
-                shared_oneway_shapes = st.shared(oneway_broadcastable_shapes())
+                shared_oneway_shapes = st.shared(hh.oneway_broadcastable_shapes())
                 left_strat = xps.arrays(
                     dtype=left_dtypes,
                     shape=shared_oneway_shapes.map(lambda S: S.result_shape),

array_api_tests/test_special_cases.py

@@ -35,10 +35,6 @@
 from . import xps
 from ._array_module import mod as xp
 from .stubs import category_to_funcs
-from .test_operators_and_elementwise_functions import (
-    oneway_broadcastable_shapes,
-    oneway_promotable_dtypes,
-)
 
 pytestmark = pytest.mark.ci
 
@@ -1281,8 +1277,8 @@ def test_binary(func_name, func, case, x1, x2, data):
 
 @pytest.mark.parametrize("iop_name, iop, case", iop_params)
 @given(
-    oneway_dtypes=oneway_promotable_dtypes(dh.float_dtypes),
-    oneway_shapes=oneway_broadcastable_shapes(),
+    oneway_dtypes=hh.oneway_promotable_dtypes(dh.float_dtypes),
+    oneway_shapes=hh.oneway_broadcastable_shapes(),
     data=st.data(),
 )
 def test_iop(iop_name, iop, case, oneway_dtypes, oneway_shapes, data):