Add writable args - #7009

Thanks very much for the PR, I've queued it for review.

Why my PR was moved to Numba 0.56RC. I need this change for my work. Is it possible to return this PR back to Numba 0.55RC?

@dmbelov Apologies for the long delay in someone reviewing your PR - it was moved to 0.56 because 0.55 RC was produced before someone could review this.

@dmbelov Many thanks for the update - I've marked this as waiting for another review.

As a small request regarding the process - could you please avoid force-pushes / rebases and instead add new commits and merge from main please? It is quite hard to tell what changed since the previous review / comments if the history is edited - many thanks in advance.

Many thanks for addressing the feedback.

As you know, API for Numpy universal and generalized universal functions does not support keywords.

It doesn't, but I don't think we need to be constrained by what the NumPy API is limited to, if we can provide a broader or easier-to-use API.

Arguments in NumPy C-API are specified by their position. Current interface for writable_args exactly mimics this C-API.

It does - I think it should be possible without too much trouble to allow arguments to be specified by name for writable_args though. I think something like:

diff --git a/numba/np/ufunc/ufuncbuilder.py b/numba/np/ufunc/ufuncbuilder.py
index b50a23d31..8b9dcabf4 100644
--- a/numba/np/ufunc/ufuncbuilder.py
+++ b/numba/np/ufunc/ufuncbuilder.py
@@ -332,7 +332,22 @@ class GUFuncBuilder(_BaseUFuncBuilder):
         self.cache = cache
         self._sigs = []
         self._cres = {}
-        self.writable_args = writable_args
+
+        args = inspect.getfullargspec(py_func).args
+        argpositions = {arg: i for i, arg in enumerate(args)}
+
+        def transform_arg(arg):
+            if isinstance(arg, int):
+                return arg
+
+            try:
+                return argpositions[arg]
+            except KeyError:
+                msg = (f"Specified writable arg {arg} not found in arg list "
+                       f"{args} for function {py_func.__qualname__}")
+                raise RuntimeError(msg)
+
+        self.writable_args = tuple([transform_arg(a) for a in writable_args])
 
     def _finalize_signature(self, cres, args, return_type):
         if not cres.objectmode and cres.signature.return_type != types.void:
diff --git a/numba/tests/npyufunc/test_update_inplace.py b/numba/tests/npyufunc/test_update_inplace.py
index 27cabdc82..dc5a0b4a4 100644
--- a/numba/tests/npyufunc/test_update_inplace.py
+++ b/numba/tests/npyufunc/test_update_inplace.py
@@ -49,6 +49,11 @@ class TestUpdateInplace(TestCase):
                              nopython=True, writable_args=(0,))(py_replace_2nd)
         self._run_test_for_gufunc(gufunc, py_replace_2nd)
 
+        # test with writable_args as strings
+        gufunc = guvectorize(['void(f8[:], f8[:])'], '(t),()',
+                             nopython=True, writable_args=('x_t',))(py_replace_2nd)
+        self._run_test_for_gufunc(gufunc, py_replace_2nd)
+
     def test_update_inplace_with_cache(self):
         # test with writable_args
         gufunc = guvectorize(['void(f8[:], f8[:])'], '(t),()',
@@ -80,6 +85,13 @@ class TestUpdateInplace(TestCase):
                              writable_args=(0, 1, 2))(py_update_3)
         self._run_test_for_gufunc(gufunc, py_update_3)
 
+        # test with writable_args as strings
+        gufunc = guvectorize(['void(f8[:], f8[:], f8[:], f8[:])'],
+                             '(t),(t),(t),()', nopython=True,
+                             writable_args=('x0_t', 'x1_t', 'x2_t'))(py_update_3)
+        self._run_test_for_gufunc(gufunc, py_update_3)
+
+
     def test_exceptions(self):
         # check that len(writable_args) <= nin
         with self.assertRaises(ValueError):

makes it possible. What do you think about the above? If it seems acceptable to you, feel free to use it in this PR - if you do, I am missing a test for the exception being raised when an argument is not found, and a doc update, which would also need to be added.

Can you merge my request to the main branch and put this change in the nearest Numba release?

This PR is in the milestone for the next release, so I do hope we can get it into the next release. I don't feel I can move this to "Ready-to-merge" yet because I'm still concerned about the API, but will also mention this PR in today's triage meeting to get the opinions of other maintainers.

The consensus from the dev meeting is that the addition of writable_args should also allow the arguments to be specified by name: draft minutes are at https://hackmd.io/n5qSqeMTR4Wv4UYkUwCZng but they will be added to the wiki over the next couple of days.

Thanks. I added changes as suggested. Please check again.

@gmarkall, do you think that this PR is ready for merging to the main branch?

@gmarkall, I just want to make sure that this PR will be a part of 0.57 release. Is there anything else that needs to be done here before you merge this PR?

@gmarkall, any updates?

@gmarkall,
Thanks for the review. I made the changes that you suggested:

• updated the documentation;
• added test of the exception raised when the user specifies a non-existent arg by name.

Please check if PR is good to be merged.

Thanks!

BFID: numba_smoketest_cpu_yaml_171

BFID: numba_smoketest_cpu_yaml_171

passed