BUG: Add further tests and fix existing npyiter cleanup bugs

2 files changed, 47 insertions, 6 deletions

diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c
index 8839d1be7..e0ec9d5cf 100644
--- a/numpy/core/src/multiarray/nditer_pywrap.c
+++ b/numpy/core/src/multiarray/nditer_pywrap.c
@@ -1134,6 +1134,9 @@ static void
 npyiter_dealloc(NewNpyArrayIterObject *self)
 {
     if (self->iter) {
+        /* Store error, so that WriteUnraisable cannot clear an existing one */
+        PyObject *exc, *val, *tb;
+        PyErr_Fetch(&exc, &val, &tb);
         if (npyiter_has_writeback(self->iter)) {
             if (PyErr_WarnEx(PyExc_RuntimeWarning,
                     "Temporary data has not been written back to one of the "
@@ -1152,10 +1155,13 @@ npyiter_dealloc(NewNpyArrayIterObject *self)
                 }
             }
         }
-        NpyIter_Deallocate(self->iter);
+        if (!NpyIter_Deallocate(self->iter)) {
+            PyErr_WriteUnraisable(Py_None);
+        }
         self->iter = NULL;
         Py_XDECREF(self->nested_child);
         self->nested_child = NULL;
+        PyErr_Restore(exc, val, tb);
     }
     Py_TYPE(self)->tp_free((PyObject*)self);
 }
@@ -2320,7 +2326,7 @@ npyiter_close(NewNpyArrayIterObject *self)
     self->iter = NULL;
     Py_XDECREF(self->nested_child);
     self->nested_child = NULL;
-    if (ret < 0) {
+    if (ret != NPY_SUCCEED) {
         return NULL;
     }
     Py_RETURN_NONE;
diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
index 536f812ee..600d0f98b 100644
--- a/numpy/core/tests/test_nditer.py
+++ b/numpy/core/tests/test_nditer.py
@@ -1963,10 +1963,45 @@ def test_iter_buffered_cast_structured_type():
     sdt2 = [('b', 'O'), ('a', 'f8')]
     a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1)
 
-    assert_raises(ValueError, lambda : (
-        nditer(a, ['buffered', 'refs_ok'], ['readwrite'],
-               casting='unsafe',
-               op_dtypes=sdt2)))
+    for intent in ["readwrite", "readonly", "writeonly"]:
+        # If the following assert fails, the place where the error is raised
+        # within nditer may change. That is fine, but it may make sense for
+        # a new (hard to design) test to replace it:
+        assert np.can_cast(a.dtype, sdt2, casting="unsafe")
+        simple_arr = np.array([1, 2], dtype="i")  # succeeds but needs clean up
+        with pytest.raises(ValueError):
+            nditer((simple_arr, a), ['buffered', 'refs_ok'], [intent, intent],
+                   casting='unsafe', op_dtypes=["f", sdt2])
+
+
+def test_buffered_cast_error_paths():
+    with pytest.raises(ValueError):
+        # The input is cast into an `S3` buffer
+        np.nditer((np.array("a", dtype="S1"),), op_dtypes=["i"],
+                  casting="unsafe", flags=["buffered"])
+
+    # The `M8[ns]` is cast into the `S3` output
+    it = np.nditer((np.array(1, dtype="i"),), op_dtypes=["S1"],
+                   op_flags=["writeonly"], casting="unsafe", flags=["buffered"])
+    with pytest.raises(ValueError):
+        with it:
+            buf = next(it)
+            buf[...] = "a"  # cannot be converted to int.
+
+    # The following gives an unraisable error, there are probably better
+    # ways to test that:
+    code = textwrap.dedent("""
+    import numpy as np
+
+    it = np.nditer((np.array(1, dtype="i"),), op_dtypes=["S1"],
+                   op_flags=["writeonly"], casting="unsafe", flags=["buffered"])
+    buf = next(it)
+    buf[...] = "a"
+    del buf, it  # Flushing only happens during deallocate right now.
+    """)
+    res = subprocess.check_output([sys.executable, "-c", code],
+                                  stderr=subprocess.STDOUT, text=True)
+    assert "ValueError" in res
 
 
 def test_iter_buffered_cast_subarray():