Merge pull request #212 from honno/fix-composite-return-hint

Fix type hints of `@st.composite` strategies

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -109,7 +109,7 @@ class OnewayPromotableDtypes(NamedTuple):
 @composite
 def oneway_promotable_dtypes(
     draw, dtypes: Sequence[DataType]
-) -> SearchStrategy[OnewayPromotableDtypes]:
+) -> 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)
@@ -127,7 +127,7 @@ class OnewayBroadcastableShapes(NamedTuple):
 
 
 @composite
-def oneway_broadcastable_shapes(draw) -> SearchStrategy[OnewayBroadcastableShapes]:
+def oneway_broadcastable_shapes(draw) -> OnewayBroadcastableShapes:
     """Return a strategy for input shapes that broadcast to result shapes."""
     result_shape = draw(shapes(min_side=1))
     input_shape = draw(

array_api_tests/test_creation_functions.py

@@ -377,7 +377,7 @@ def test_eye(n_rows, n_cols, kw):
 
 
 @st.composite
-def full_fill_values(draw) -> st.SearchStrategy[Union[bool, int, float, complex]]:
+def full_fill_values(draw) -> Union[bool, int, float, complex]:
     kw = draw(
         st.shared(hh.kwargs(dtype=st.none() | xps.scalar_dtypes()), key="full_kw")
     )