ENH: missingdata: Add maskna= parameter to np.linspace and np.logspace

2 files changed, 26 insertions, 7 deletions

diff --git a/numpy/core/function_base.py b/numpy/core/function_base.py
index b2f9dc70c..3e919c761 100644
--- a/numpy/core/function_base.py
+++ b/numpy/core/function_base.py
@@ -3,7 +3,7 @@ __all__ = ['logspace', 'linspace']
 import numeric as _nx
 from numeric import array
 
-def linspace(start, stop, num=50, endpoint=True, retstep=False):
+def linspace(start, stop, num=50, endpoint=True, retstep=False, maskna=False):
     """
     Return evenly spaced numbers over a specified interval.
 
@@ -29,6 +29,8 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False):
     retstep : bool, optional
         If True, return (`samples`, `step`), where `step` is the spacing
         between samples.
+    maskna : boolean
+        If this is true, the returned array will have an NA mask.
 
     Returns
     -------
@@ -73,22 +75,22 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False):
     """
     num = int(num)
     if num <= 0:
-        return array([], float)
+        return array([], float, maskna=maskna)
     if endpoint:
         if num == 1:
-            return array([float(start)])
+            return array([float(start)], maskna=maskna)
         step = (stop-start)/float((num-1))
-        y = _nx.arange(0, num) * step + start
+        y = _nx.arange(0, num, maskna=maskna) * step + start
         y[-1] = stop
     else:
         step = (stop-start)/float(num)
-        y = _nx.arange(0, num) * step + start
+        y = _nx.arange(0, num, maskna=maskna) * step + start
     if retstep:
         return y, step
     else:
         return y
 
-def logspace(start,stop,num=50,endpoint=True,base=10.0):
+def logspace(start,stop,num=50,endpoint=True,base=10.0, maskna=False):
     """
     Return numbers spaced evenly on a log scale.
 
@@ -114,6 +116,8 @@ def logspace(start,stop,num=50,endpoint=True,base=10.0):
         The base of the log space. The step size between the elements in
         ``ln(samples) / ln(base)`` (or ``log_base(samples)``) is uniform.
         Default is 10.0.
+    maskna : boolean
+        If this is true, the returned array will have an NA mask.
 
     Returns
     -------
@@ -162,6 +166,6 @@ def logspace(start,stop,num=50,endpoint=True,base=10.0):
     >>> plt.show()
 
     """
-    y = linspace(start,stop,num=num,endpoint=endpoint)
+    y = linspace(start,stop,num=num,endpoint=endpoint,maskna=maskna)
     return _nx.power(base,y)
 
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index 8fd3fb174..35165cc24 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -1157,5 +1157,20 @@ def test_array_maskna_var_std():
     res = np.std(a, axis=1, skipna=True)
     assert_array_almost_equal(res, [1.0, 0.81649658092772603])
 
+def test_array_maskna_linspace_logspace():
+    # np.linspace, np.logspace
+
+    a = np.linspace(2.0, 3.0, num=5)
+    b = np.linspace(2.0, 3.0, num=5, maskna=True)
+    assert_equal(a, b)
+    assert_(not a.flags.maskna)
+    assert_(b.flags.maskna)
+
+    a = np.logspace(2.0, 3.0, num=4)
+    b = np.logspace(2.0, 3.0, num=4, maskna=True)
+    assert_equal(a, b)
+    assert_(not a.flags.maskna)
+    assert_(b.flags.maskna)
+
 if __name__ == "__main__":
     run_module_suite()