FIX: random.draw: Add =None

Fix #704

Author: oyamad

quantecon/random/tests/test_utilities.py

@@ -108,3 +108,14 @@ def test_lln(self):
             pmf_computed = hist * np.diff(bin_edges)
             atol = 1e-2
             assert_allclose(pmf_computed, self.pmf, atol=atol)
+
+
+@njit
+def draw_jitted_w_o_size(n):
+    cdf = np.linspace(1/n, 1, n)
+    return draw(cdf)
+
+
+def test_draw_jitted_w_o_size():
+    n = 3
+    assert_(draw_jitted_w_o_size(n) in range(n))

quantecon/random/utilities.py

@@ -211,16 +211,16 @@ def draw(cdf, size=None):
 
 # Overload for the `draw` function
 @overload(draw)
-def ol_draw(cdf, size):
+def ol_draw(cdf, size=None):
     if isinstance(size, types.Integer):
-        def draw_impl(cdf, size):
+        def draw_impl(cdf, size=None):
             rs = np.random.random(size)
             out = np.empty(size, dtype=np.int_)
             for i in range(size):
                 out[i] = searchsorted(cdf, rs[i])
             return out
     else:
-        def draw_impl(cdf, size):
+        def draw_impl(cdf, size=None):
             r = np.random.random()
             return searchsorted(cdf, r)
     return draw_impl