Merge 8ef9e04 into 1871d6f

bashtage · Apr 12, 2019 · efb61dc · efb61dc
2 parents 1871d6f + 8ef9e04
commit efb61dc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/randomgen/generator.pyx b/randomgen/generator.pyx
@@ -3801,11 +3801,12 @@ cdef class RandomGenerator:
         on = <np.ndarray>np.PyArray_FROM_OTF(n, np.NPY_INT64, np.NPY_ALIGNED)
         parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
         pix = <double*>np.PyArray_DATA(parr)
-
+        check_array_constraint(parr, 'pvals', CONS_BOUNDED_0_1)
         if kahan_sum(pix, d-1) > (1.0 + 1e-12):
             raise ValueError("sum(pvals[:-1]) > 1.0")
 
         if np.PyArray_NDIM(on) != 0: # vector
+            check_array_constraint(on, 'n', CONS_NON_NEGATIVE)
             if size is None:
                 it = np.PyArray_MultiIterNew1(on)
             else:
@@ -3839,6 +3840,7 @@ cdef class RandomGenerator:
         mnix = <int64_t*>np.PyArray_DATA(mnarr)
         sz = np.PyArray_SIZE(mnarr)
         ni = n
+        check_constraint(ni, 'n', CONS_NON_NEGATIVE)
         offset = 0
         with self.lock, nogil:
             for i in range(sz // d):

diff --git a/randomgen/tests/test_generator_mt19937.py b/randomgen/tests/test_generator_mt19937.py
@@ -90,6 +90,11 @@ def test_size(self):
 
     def test_invalid_prob(self):
         assert_raises(ValueError, random.multinomial, 100, [1.1, 0.2])
+        assert_raises(ValueError, random.multinomial, 100, [-.1, 0.9])
+
+    def test_invalid_n(self):
+        assert_raises(ValueError, random.multinomial, -1, [0.8, 0.2])
+        assert_raises(ValueError, random.multinomial, [-1] * 10, [0.8, 0.2])
 
 
 class TestSetState(object):