writable_args can be specified as tuple of argument positions, argume…

…nt names, or mixture of both.
sklam · Mar 28, 2023 · May 9, 2021 · May 9, 2021 · May 9, 2021 · May 10, 2021
commit b0c435ef29525b05ba9d1502bd81f8a48a1f0b48
diff --git a/docs/source/user/vectorize.rst b/docs/source/user/vectorize.rst
@@ -302,7 +302,8 @@ casted array was mutated instead.
    ``writable_args=(n, m)``.
 
 To solve this problem, one needs to tell the GUFunc engine that 0-th
-argument is writable. This can be achieved by passing ``writable_args=(0,)``
+argument is writable. This can be achieved by passing argument position
+``writable_args=(0,)`` or argument name ``writable_args=('invals',)``
 to ``@guvectorize``.  Now, the code above works as expected::
 
    @guvectorize([(float64[:], float64[:])], '()->()', writable_args=(0,))

diff --git a/numba/np/ufunc/ufuncbuilder.py b/numba/np/ufunc/ufuncbuilder.py
@@ -332,7 +332,9 @@ def __init__(self, py_func, signature, identity=None, cache=False,
         self.cache = cache
         self._sigs = []
         self._cres = {}
-        self.writable_args = writable_args
+
+        transform_arg = _get_transform_arg(py_func)
+        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:
@@ -394,3 +396,22 @@ def build(self, cres):
                 ty = a
             dtypenums.append(as_dtype(ty).num)
         return dtypenums, ptr, env
+
+
+def _get_transform_arg(py_func):
+    """Return function that transform arg into index"""
+    args = inspect.getfullargspec(py_func).args
+    pos_by_arg = {arg: i for i, arg in enumerate(args)}
+
+    def transform_arg(arg):
+        if isinstance(arg, int):
+            return arg
+
+        try:
+            return pos_by_arg[arg]
+        except KeyError:
+            msg = (f"Specified writable arg {arg} not found in arg list "
+                   f"{args} for function {py_func.__qualname__}")
+            raise RuntimeError(msg)
+
+    return transform_arg
diff --git a/numba/tests/npyufunc/test_update_inplace.py b/numba/tests/npyufunc/test_update_inplace.py
@@ -49,6 +49,11 @@ def test_update_inplace(self):
                              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,12 @@ def test_update_inplace_3(self):
                              writable_args=(0, 1, 2))(py_update_3)
         self._run_test_for_gufunc(gufunc, py_update_3)
 
+        # test with writable_args as mix of strings and ints
+        gufunc = guvectorize(['void(f8[:], f8[:], f8[:], f8[:])'],
+                             '(t),(t),(t),()', nopython=True,
+                             writable_args=('x0_t', 'x1_t', 2))(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):