diff options
| -rw-r--r-- | numpy/core/src/multiarray/dtype_transfer.c | 1 | ||||
| -rw-r--r-- | numpy/core/tests/test_nditer.py | 31 |
2 files changed, 32 insertions, 0 deletions
diff --git a/numpy/core/src/multiarray/dtype_transfer.c b/numpy/core/src/multiarray/dtype_transfer.c index 9b8e5f32f..fa17caced 100644 --- a/numpy/core/src/multiarray/dtype_transfer.c +++ b/numpy/core/src/multiarray/dtype_transfer.c @@ -4625,6 +4625,7 @@ PyArray_GetMaskedDTypeTransferFunction(int aligned, /* TODO: Special case some important cases so they're fast */ /* Fall back to wrapping a non-masked transfer function */ + assert(dst_dtype != NULL); if (PyArray_GetDTypeTransferFunction(aligned, src_stride, dst_stride, src_dtype, dst_dtype, diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py index 4271d2d96..411095199 100644 --- a/numpy/core/tests/test_nditer.py +++ b/numpy/core/tests/test_nditer.py @@ -2746,6 +2746,37 @@ def test_iter_writemasked(): # were copied back assert_equal(a, [3, 3, 2.5]) +def test_iter_writemasked_decref(): + # force casting (to make it interesting) by using a structured dtype. + arr = np.arange(10000).astype(">i,O") + original = arr.copy() + mask = np.random.randint(0, 2, size=10000).astype(bool) + + it = np.nditer([arr, mask], ['buffered', "refs_ok"], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask']], + op_dtypes=["<i,O", "?"]) + singleton = object() + if HAS_REFCOUNT: + count = sys.getrefcount(singleton) + for buf, mask_buf in it: + buf[...] = (3, singleton) + + del buf, mask_buf, it # delete everything to ensure corrrect cleanup + + if HAS_REFCOUNT: + # The buffer would have included additional items, they must be + # cleared correctly: + assert sys.getrefcount(singleton) - count == np.count_nonzero(mask) + + assert_array_equal(arr[~mask], original[~mask]) + assert (arr[mask] == np.array((3, singleton), arr.dtype)).all() + del arr + + if HAS_REFCOUNT: + assert sys.getrefcount(singleton) == count + + def test_iter_non_writable_attribute_deletion(): it = np.nditer(np.ones(2)) attr = ["value", "shape", "operands", "itviews", "has_delayed_bufalloc", |