More thorough test_full()

Also implements `DType.kind` to raise `NotImplementedError`

Author: honno

torch_np/_dtypes.py

@@ -55,6 +55,10 @@ def name(self):
     def type(self):
         return self._scalar_type
 
+    @property
+    def kind(self):
+        raise NotImplementedError
+
     @property
     def typecode(self):
         return self._scalar_type.typecode

torch_np/tests/test_xps.py

@@ -25,34 +25,40 @@
     xps = make_strategies_namespace(np, api_version="2022.12")
 
 
+default_dtypes = [np.bool, np.int64, np.float64, np.complex128]
+kind_to_strat = {
+    "b": xps.boolean_dtypes(),
+    "i": xps.integer_dtypes(),
+    "u": xps.unsigned_integer_dtypes(sizes=8),
+    "f": xps.floating_dtypes(),
+    "c": xps.complex_dtypes(),
+}
+scalar_dtype_strat = st.one_of(kind_to_strat.values()).map(np.dtype)
+
+
 @given(shape=xps.array_shapes(), data=st.data())
 def test_full(shape, data):
     if data.draw(st.booleans(), label="pass kwargs?"):
-        kw = {}
-    else:
-        dtype = data.draw(st.none() | xps.scalar_dtypes(), label="dtype")
+        dtype = data.draw(st.none() | scalar_dtype_strat, label="dtype")
         kw = {"dtype": dtype}
-    _dtype = kw.get("dtype", None) or data.draw(
-        st.sampled_from([np.bool, np.int64, np.float64, np.complex128]), label="_dtype"
-    )
+    else:
+        kw = {}
+    _dtype = kw.get("dtype", None) or data.draw(scalar_dtype_strat, label="_dtype")
     values_strat = xps.from_dtype(_dtype)
-    fill_value = data.draw(
-        values_strat | values_strat.map(lambda v: np.asarray(v, dtype=_dtype)),
-        label="fill_value",
-    )
-    out = np.full(shape, fill_value, **kw)
-    if kw.get("dtype", None) is None and not isinstance(fill_value, np.ndarray):
-        if isinstance(fill_value, bool):
-            assert out.dtype == np.bool
-        elif isinstance(fill_value, int):
-            assert out.dtype == np.int64
-        elif isinstance(fill_value, float):
-            assert out.dtype == np.float64
+    if _dtype not in default_dtypes or data.draw(
+        st.booleans(), label="fill_value is array?"
+    ):
+        if specified_dtype := kw.get("dtype", None):
+            kind = specified_dtype.name[0]
+            values_dtypes_strat = kind_to_strat[kind]
         else:
-            assert isinstance(fill_value, complex)  # sanity check
-            assert out.dtype == np.complex128
-    else:
-        assert out.dtype == _dtype
+            values_dtypes_strat = st.just(_dtype)
+        values_strat = values_dtypes_strat.flatmap(
+            lambda d: values_strat.map(lambda v: np.asarray(v, dtype=d))
+        )
+    fill_value = data.draw(values_strat, label="fill_value")
+    out = np.full(shape, fill_value, **kw)
+    assert out.dtype == _dtype
     assert out.shape == shape
     if cmath.isnan(fill_value):
         assert np.isnan(out).all()