ENH: core: Add tests for copyto function with new array_assign_array features

2 files changed, 75 insertions, 1 deletions

diff --git a/numpy/core/src/multiarray/array_assign_array.c b/numpy/core/src/multiarray/array_assign_array.c
index bdc262e17..96bfd1f45 100644
--- a/numpy/core/src/multiarray/array_assign_array.c
+++ b/numpy/core/src/multiarray/array_assign_array.c
@@ -319,7 +319,7 @@ raw_array_wheremasked_assign_array_preservena(int ndim, npy_intp *shape,
         /* Process the innermost dimension a buffer size at a time */
         count = shape_it[0];
         dst_d = dst_data;
-        src_d = dst_data;
+        src_d = src_data;
         maskna_d = maskna_data;
         wheremask_d = wheremask_data;
         do {
diff --git a/numpy/core/tests/test_api.py b/numpy/core/tests/test_api.py
index 8ba1b4806..f98c8c8c1 100644
--- a/numpy/core/tests/test_api.py
+++ b/numpy/core/tests/test_api.py
@@ -167,5 +167,79 @@ def test_copyto():
     # 'dst' must be an array
     assert_raises(TypeError, np.copyto, [1,2,3], [2,3,4])
 
+def test_copyto_maskna():
+    a_orig = np.zeros((2,3), dtype='f8')
+    a = a_orig.view(maskna=True)
+
+    # Simple copy to from non-masked to NA-masked
+    a[...] = np.NA
+    np.copyto(a, np.arange(6).reshape(2,3))
+    assert_equal(a, [[0,1,2],[3,4,5]])
+    a[...] = np.NA
+    np.copyto(a.T, np.arange(6).reshape(3,2) + 1)
+    assert_equal(a, [[1,3,5],[2,4,6]])
+
+    # Simple copy to from NA-masked to NA-masked
+    a[...] = np.NA
+    a[1,2] = 12
+    tmp = np.arange(6, maskna=True).reshape(2,3)
+    tmp[0,1] = np.NA
+    tmp[1,2] = np.NA
+    np.copyto(a, tmp)
+    assert_equal(a_orig, [[0,3,2],[3,4,12]])
+    assert_equal(np.isna(a), [[0,1,0],[0,0,1]])
+
+    # Where-masked copy from non-masked to NA-masked
+    a[...] = np.NA
+    a[0,2] = 6
+    mask = np.array([[0,0,1],[0,1,0]], dtype='?')
+    tmp = np.arange(6).reshape(2,3) + 1
+    np.copyto(a, tmp, where=mask)
+    assert_equal(a_orig, [[0,3,3],[3,5,12]])
+    assert_equal(np.isna(a), ~mask)
+
+    # Where-masked copy from NA-masked to NA-masked
+    a[1,2] = 12
+    mask = np.array([[0,1,1],[0,0,1]], dtype='?')
+    tmp = np.arange(6, maskna=True).reshape(2,3) + 3
+    tmp[0,0] = np.NA
+    tmp[0,1] = np.NA
+    tmp[1,2] = np.NA
+    np.copyto(a, tmp, where=mask)
+    assert_equal(a_orig, [[0,3,5],[3,5,12]])
+    assert_equal(np.isna(a), [[1,1,0],[1,0,1]])
+
+    # Preserve-NA copy from non-masked to NA-masked
+    np.copyto(a, np.arange(6).reshape(2,3), preservena=True)
+    assert_equal(a_orig, [[0,3,2],[3,4,12]])
+    assert_equal(np.isna(a), [[1,1,0],[1,0,1]])
+
+    # Preserve-NA copy from NA-masked to NA-masked
+    tmp = np.arange(6, maskna=True).reshape(2,3) + 1
+    tmp[0,0] = np.NA
+    tmp[1,1] = np.NA
+    np.copyto(a, tmp, preservena=True)
+    assert_equal(a_orig, [[0,3,3],[3,4,12]])
+    assert_equal(np.isna(a), [[1,1,0],[1,1,1]])
+
+    # Where-masked preserve-NA copy from non-masked to NA-masked
+    tmp = np.arange(6).reshape(2,3) + 3
+    a[1,2] = 12
+    mask = np.array([[0,1,1],[0,1,0]], dtype='?')
+    np.copyto(a, tmp, where=mask, preservena=True)
+    assert_equal(a_orig, [[0,3,5],[3,4,12]])
+    assert_equal(np.isna(a), [[1,1,0],[1,1,0]])
+
+    # Where-masked preserve-NA copy from NA-masked to NA-masked
+    a[0,0] = 0
+    mask = np.array([[0,1,1],[1,0,1]], dtype='?')
+    tmp = np.arange(6, maskna=True).reshape(2,3) + 1
+    tmp[1,0] = np.NA
+    tmp[1,2] = np.NA
+    np.copyto(a, tmp, where=mask, preservena=True)
+    assert_equal(a_orig, [[0,3,3],[3,4,12]])
+    assert_equal(np.isna(a), [[0,1,0],[1,1,1]])
+
+
 if __name__ == "__main__":
     run_module_suite()