Filter undefined dtypes in mutually_promotable_dtypes

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -2,7 +2,7 @@
 from functools import reduce
 from math import sqrt
 from operator import mul
-from typing import Any, List, NamedTuple, Optional, Tuple
+from typing import Any, List, NamedTuple, Optional, Tuple, Sequence
 
 from hypothesis import assume
 from hypothesis.strategies import (SearchStrategy, booleans, composite, floats,
@@ -68,19 +68,14 @@ def _dtypes_sorter(dtype_pair: Tuple[DataType, DataType]):
 
 promotable_dtypes: List[Tuple[DataType, DataType]] = sorted(dh.promotion_table.keys(), key=_dtypes_sorter)
 
-if FILTER_UNDEFINED_DTYPES:
-    promotable_dtypes = [
-        (i, j) for i, j in promotable_dtypes
-        if not isinstance(i, _UndefinedStub)
-        and not isinstance(j, _UndefinedStub)
-    ]
-
-
 def mutually_promotable_dtypes(
     max_size: Optional[int] = 2,
     *,
-    dtypes: Tuple[DataType, ...] = dh.all_dtypes,
+    dtypes: Sequence[DataType] = dh.all_dtypes,
 ) -> SearchStrategy[Tuple[DataType, ...]]:
+    if FILTER_UNDEFINED_DTYPES:
+        dtypes = [d for d in dtypes if not isinstance(d, _UndefinedStub)]
+        assert len(dtypes) > 0, "all dtypes undefined"  # sanity check
     if max_size == 2:
         return sampled_from(
             [(i, j) for i, j in promotable_dtypes if i in dtypes and j in dtypes]
@@ -347,7 +342,7 @@ def multiaxis_indices(draw, shapes):
 
 
 def two_mutual_arrays(
-    dtypes: Tuple[DataType, ...] = dh.all_dtypes,
+    dtypes: Sequence[DataType] = dh.all_dtypes,
     two_shapes: SearchStrategy[Tuple[Shape, Shape]] = two_mutually_broadcastable_shapes,
 ) -> SearchStrategy:
     mutual_dtypes = shared(mutually_promotable_dtypes(dtypes=dtypes))

array_api_tests/meta/test_hypothesis_helpers.py

@@ -15,15 +15,29 @@
 pytestmark = [pytest.mark.skipif(UNDEFINED_DTYPES, reason="undefined dtypes")]
 
 @given(hh.mutually_promotable_dtypes(dtypes=dh.float_dtypes))
-def test_mutually_promotable_dtypes(pairs):
-    assert pairs in (
+def test_mutually_promotable_dtypes(pair):
+    assert pair in (
         (xp.float32, xp.float32),
         (xp.float32, xp.float64),
         (xp.float64, xp.float32),
         (xp.float64, xp.float64),
     )
 
 
+@given(
+    hh.mutually_promotable_dtypes(
+        dtypes=[xp.uint8, _UndefinedStub("uint16"), xp.uint32]
+    )
+)
+def test_partial_mutually_promotable_dtypes(pair):
+    assert pair in (
+        (xp.uint8, xp.uint8),
+        (xp.uint8, xp.uint32),
+        (xp.uint32, xp.uint8),
+        (xp.uint32, xp.uint32),
+    )
+
+
 def valid_shape(shape) -> bool:
     return (
         all(isinstance(side, int) for side in shape)