Add further tests to RandomState and Generator

Quansight-Labs · Aug 27, 2020 · ba6a968 · ba6a968
1 parent 9da7501
commit ba6a968
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/unumpy/sparse_backend.py b/unumpy/sparse_backend.py
@@ -2,7 +2,7 @@
 import sparse
 from uarray import Dispatchable, wrap_single_convertor
 from unumpy import ufunc, ufunc_list, ndarray, dtype
-from unumpy.random import RandomState
+from unumpy.random import RandomState, Generator
 import unumpy
 import functools
 
@@ -28,7 +28,8 @@ def array(x, *args, **kwargs):
     ndarray: sparse.SparseArray,
     dtype: np.dtype,
     ufunc: np.ufunc,
-    RandomState: np.random.mtrand.RandomState,
+    RandomState: np.random.RandomState,
+    Generator: np.random.Generator,
 }
 
 

diff --git a/unumpy/tests/test_numpy.py b/unumpy/tests/test_numpy.py
@@ -804,26 +804,38 @@ def test_class_overriding():
         assert isinstance(onp.dtype("float64"), np.dtype)
         assert np.dtype("float64") == onp.float64
         assert isinstance(np.dtype("float64"), onp.dtype)
+        assert isinstance(onp.random.RandomState(), np.random.RandomState)
+        assert isinstance(onp.random.Generator(onp.random.PCG64()), np.random.Generator)
         assert issubclass(onp.ufunc, np.ufunc)
+        assert issubclass(onp.random.RandomState, np.random.RandomState)
+        assert issubclass(onp.random.Generator, np.random.Generator)
 
     with ua.set_backend(DaskBackend(), coerce=True):
         assert isinstance(da.add, np.ufunc)
         assert isinstance(onp.dtype("float64"), np.dtype)
         assert np.dtype("float64") == onp.float64
         assert isinstance(np.dtype("float64"), onp.dtype)
+        assert isinstance(da.random.RandomState(), np.random.RandomState)
         assert issubclass(da.ufunc.ufunc, np.ufunc)
+        assert issubclass(da.random.RandomState, np.random.RandomState)
 
     with ua.set_backend(SparseBackend, coerce=True):
         assert isinstance(onp.add, np.ufunc)
         assert isinstance(onp.dtype("float64"), np.dtype)
         assert np.dtype("float64") == onp.float64
         assert isinstance(np.dtype("float64"), onp.dtype)
+        assert isinstance(onp.random.RandomState(), np.random.RandomState)
+        assert isinstance(onp.random.Generator(onp.random.PCG64()), np.random.Generator)
         assert issubclass(onp.ufunc, np.ufunc)
+        assert issubclass(onp.random.RandomState, np.random.RandomState)
+        assert issubclass(onp.random.Generator, np.random.Generator)
 
     if hasattr(CupyBackend, "__ua_function__"):
         with ua.set_backend(CupyBackend, coerce=True):
             assert isinstance(cp.add, np.ufunc)
             assert isinstance(cp.dtype("float64"), np.dtype)
             assert np.dtype("float64") == cp.float64
             assert isinstance(np.dtype("float64"), cp.dtype)
+            assert isinstance(cp.random.RandomState(), np.random.RandomState)
             assert issubclass(cp.ufunc, np.ufunc)
+            assert issubclass(cp.random.RandomState, np.random.RandomState)