ENH: missingdata: Add maskna= flag to np.eye constructor

3 files changed, 24 insertions, 11 deletions

diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py
index 3c05cf71d..33de34e1a 100644
--- a/numpy/core/_methods.py
+++ b/numpy/core/_methods.py
@@ -98,7 +98,7 @@ def _std(a, axis=None, dtype=None, out=None, ddof=0,
                                 skipna=skipna, keepdims=keepdims)
 
     if isinstance(ret, mu.ndarray):
-        um.sqrt(ret, out=ret)
+        ret = um.sqrt(ret, out=ret)
     else:
         ret = um.sqrt(ret)
 
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index 638c52ba6..77ab22b29 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -1338,5 +1338,22 @@ def test_array_maskna_linspace_logspace():
     assert_(not a.flags.maskna)
     assert_(b.flags.maskna)
 
+def test_array_maskna_eye():
+    # np.eye
+
+    # By default there should be no NA mask
+    a = np.eye(3)
+    assert_(not a.flags.maskna)
+
+    a = np.eye(3, maskna=True)
+    assert_(a.flags.maskna)
+    assert_(a.flags.ownmaskna)
+    assert_equal(a, np.eye(3))
+
+    a = np.eye(3, k=2, maskna=True)
+    assert_(a.flags.maskna)
+    assert_(a.flags.ownmaskna)
+    assert_equal(a, np.eye(3, k=2))
+
 if __name__ == "__main__":
     run_module_suite()
diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py
index 9b3b50d04..12bba99a6 100644
--- a/numpy/lib/twodim_base.py
+++ b/numpy/lib/twodim_base.py
@@ -9,7 +9,7 @@ __all__ = ['diag','diagflat','eye','fliplr','flipud','rot90','tri','triu',
 
 from numpy.core.numeric import asanyarray, equal, subtract, arange, \
      zeros, greater_equal, multiply, ones, asarray, alltrue, where, \
-     empty
+     empty, diagonal
 
 def fliplr(m):
     """
@@ -166,7 +166,7 @@ def rot90(m, k=1):
         # k == 3
         return fliplr(m.swapaxes(0,1))
 
-def eye(N, M=None, k=0, dtype=float):
+def eye(N, M=None, k=0, dtype=float, maskna=False):
     """
     Return a 2-D array with ones on the diagonal and zeros elsewhere.
 
@@ -182,6 +182,8 @@ def eye(N, M=None, k=0, dtype=float):
       to a lower diagonal.
     dtype : data-type, optional
       Data-type of the returned array.
+    maskna : boolean
+      If this is true, the returned array will have an NA mask.
 
     Returns
     -------
@@ -207,14 +209,8 @@ def eye(N, M=None, k=0, dtype=float):
     """
     if M is None:
         M = N
-    m = zeros((N, M), dtype=dtype)
-    if k >= M:
-        return m
-    if k >= 0:
-        i = k
-    else:
-        i = (-k) * M
-    m[:M-k].flat[i::M+1] = 1
+    m = zeros((N, M), dtype=dtype, maskna=maskna)
+    diagonal(m, k)[...] = 1
     return m
 
 def diag(v, k=0):