[Numpy-discussion] make nditer also a context manager

Matti Picus matti.picus at gmail.com
Tue Jul 4 00:15:24 EDT 2017

When an nditer uses certain op_flags[0], like updateifcopy or readwrite 
or copy, the operands (which are ndarray views into the original data) 
must use the UPDATEIFCOPY flag. The meaning of this flag is to allocate 
temporary memory to hold the modified data and make the original data 
readonly. When the caller is finished with the nditer, the temporary 
memory is written back to to the original data and the original data's 
readwrite status is restored. The trigger to resolve the temporary data 
is currently via the deallocate function of the nditer, thus the call 
becomes something like

i = nditer(a, op_flags=<flags>)
#do something with i
i = None # trigger writing data from i to a

This IMO violates the "explicit is better" philosophy, and has the 
additional disadvantage of relying on refcounting semantics to trigger 
the data resolution, which does not work on PyPy.

I have a pending pull request[1] to add a private numpy API function 
PyArray_ResolveUpdateIfCopy, which allows triggering the data resolution 
in a more explicit way. The new API function is added at the end of the 
functions like take or put with non-contiguous arguments, explicit tests 
have also been added . The only user-facing effect is to allow using an 
nditer as a context manager, so the lines above would become

with nditer(a, op_flags=<flags>) as i:
     #do something with i
# data is written back when exiting

The pull request passes all tests on CPython. The last commit makes the 
use of a nditer context manager mandatory on PyPy if the UPDATEIFCOPY 
semantics are triggered, while allowing existing code to function 
without a warning on CPython.

Note that np.nested_iters[2] is problematic, in that it currently 
returns a tuple of nidters, which AFAICT cannot be made into a context 
manager, so __enter__ and __exit__ must be called manually for the nditers.

At some future point we could decide to deprecate the non-context 
managed use on CPython as well, following a cycle of first issuing a 
deprecation warning for a few release versions.

Any thoughts? Does making nditer a context manager make sense? How 
widespread is use of nested_iters in live code (it does not appear in 
the official NumPy documentation)

[0] https://docs.scipy.org/doc/numpy/reference/generated/numpy.nditer.html
[1] https://github.com/numpy/numpy/pull/9269

More information about the NumPy-Discussion mailing list