patch 9.1.0547: No way to get the arity of a Vim function

Problem:  No way to get the arity of a Vim function
          (Austin Ziegler)
Solution: Enhance get() Vim script function to return the function
          argument info using get(func, "arity") (LemonBoy)

fixes: #15097
closes: #15109

Signed-off-by: LemonBoy <thatlemon@gmail.com>
Signed-off-by: Christian Brabandt <cb@256bit.org>

src/evalfunc.c

@@ -5134,6 +5134,36 @@ f_get(typval_T *argvars, typval_T *rettv)
 			list_append_tv(rettv->vval.v_list, &pt->pt_argv[i]);
 		}
 	    }
+	    else if (STRCMP(what, "arity") == 0)
+	    {
+		int required = 0, optional = 0, varargs = FALSE;
+		char_u *name = partial_name(pt);
+
+		get_func_arity(name, &required, &optional, &varargs);
+
+		rettv->v_type = VAR_DICT;
+		if (rettv_dict_alloc(rettv) == OK)
+		{
+		    dict_T *dict = rettv->vval.v_dict;
+
+		    // Take into account the arguments of the partial, if any.
+		    // Note that it is possible to supply more arguments than the function
+		    // accepts.
+		    if (pt->pt_argc >= required + optional)
+			required = optional = 0;
+		    else if (pt->pt_argc > required)
+		    {
+			optional -= pt->pt_argc - required;
+			required = 0;
+		    }
+		    else
+			required -= pt->pt_argc;
+
+		    dict_add_number(dict, "required", required);
+		    dict_add_number(dict, "optional", optional);
+		    dict_add_bool(dict, "varargs", varargs);
+		}
+	    }
 	    else
 		semsg(_(e_invalid_argument_str), what);
 

src/proto/userfunc.pro

@@ -95,4 +95,5 @@ int set_ref_in_call_stack(int copyID);
 int set_ref_in_functions(int copyID);
 int set_ref_in_func_args(int copyID);
 int set_ref_in_func(char_u *name, ufunc_T *fp_in, int copyID);
+int get_func_arity(char_u *name, int *required, int *optional, int *varargs);
 /* vim: set ft=c : */

src/testdir/test_getvar.vim

@@ -142,20 +142,28 @@ func Test_get_func()
   let l:F = function('tr')
   call assert_equal('tr', get(l:F, 'name'))
   call assert_equal(l:F, get(l:F, 'func'))
+  call assert_equal({'required': 3, 'optional': 0, 'varargs': v:false},
+      \ get(l:F, 'arity'))
 
   let Fb_func = function('s:FooBar')
   call assert_match('<SNR>\d\+_FooBar', get(Fb_func, 'name'))
+  call assert_equal({'required': 0, 'optional': 0, 'varargs': v:false},
+      \ get(Fb_func, 'arity'))
   let Fb_ref = funcref('s:FooBar')
   call assert_match('<SNR>\d\+_FooBar', get(Fb_ref, 'name'))
+  call assert_equal({'required': 0, 'optional': 0, 'varargs': v:false},
+      \ get(Fb_ref, 'arity'))
 
   call assert_equal({'func has': 'no dict'}, get(l:F, 'dict', {'func has': 'no dict'}))
   call assert_equal(0, get(l:F, 'dict'))
   call assert_equal([], get(l:F, 'args'))
+
   let NF = test_null_function()
   call assert_equal('', get(NF, 'name'))
   call assert_equal(NF, get(NF, 'func'))
   call assert_equal(0, get(NF, 'dict'))
   call assert_equal([], get(NF, 'args'))
+  call assert_equal({'required': 0, 'optional': 0, 'varargs': v:false}, get(NF, 'arity'))
 endfunc
 
 " get({partial}, {what} [, {default}]) - in test_partial.vim

src/testdir/test_partial.vim

@@ -311,6 +311,11 @@ func Test_auto_partial_rebind()
 endfunc
 
 func Test_get_partial_items()
+  func s:Qux(x, y, z=3, w=1, ...)
+  endfunc
+  func s:Qux1(x, y)
+  endfunc
+
   let dict = {'name': 'hello'}
   let args = ["foo", "bar"]
   let Func = function('MyDictFunc')
@@ -331,6 +336,23 @@ func Test_get_partial_items()
   let dict = {'partial has': 'no dict'}
   call assert_equal(dict, get(P, 'dict', dict))
   call assert_equal(0, get(l:P, 'dict'))
+
+  call assert_equal({'required': 2, 'optional': 2, 'varargs': v:true},
+      \ get(funcref('s:Qux', []), 'arity'))
+  call assert_equal({'required': 1, 'optional': 2, 'varargs': v:true},
+      \ get(funcref('s:Qux', [1]), 'arity'))
+  call assert_equal({'required': 0, 'optional': 2, 'varargs': v:true},
+      \ get(funcref('s:Qux', [1, 2]), 'arity'))
+  call assert_equal({'required': 0, 'optional': 1, 'varargs': v:true},
+      \ get(funcref('s:Qux', [1, 2, 3]), 'arity'))
+  call assert_equal({'required': 0, 'optional': 0, 'varargs': v:true},
+      \ get(funcref('s:Qux', [1, 2, 3, 4]), 'arity'))
+  " More args than expected is not an error
+  call assert_equal({'required': 0, 'optional': 0, 'varargs': v:false},
+      \ get(funcref('s:Qux1', [1, 2, 3, 4]), 'arity'))
+
+  delfunc s:Qux
+  delfunc s:Qux1
 endfunc
 
 func Test_compare_partials()

src/userfunc.c

@@ -5503,6 +5503,47 @@ ex_function(exarg_T *eap)
     ga_clear_strings(&lines_to_free);
 }
 
+    int
+get_func_arity(char_u *name, int *required, int *optional, int *varargs)
+{
+    ufunc_T	*ufunc = NULL;
+    int		argcount = 0;
+    int		min_argcount = 0;
+    int		idx;
+
+    idx = find_internal_func(name);
+    if (idx >= 0)
+    {
+	internal_func_get_argcount(idx, &argcount, &min_argcount);
+	*varargs = FALSE;
+    }
+    else
+    {
+	char_u	fname_buf[FLEN_FIXED + 1];
+	char_u	*tofree = NULL;
+	funcerror_T error = FCERR_NONE;
+	char_u	*fname;
+
+	// May need to translate <SNR>123_ to K_SNR.
+	fname = fname_trans_sid(name, fname_buf, &tofree, &error);
+	if (error == FCERR_NONE)
+	    ufunc = find_func(fname, FALSE);
+	vim_free(tofree);
+
+	if (ufunc == NULL)
+	    return FAIL;
+
+	argcount = ufunc->uf_args.ga_len;
+	min_argcount = ufunc->uf_args.ga_len - ufunc->uf_def_args.ga_len;
+	*varargs = has_varargs(ufunc);
+    }
+
+    *required = min_argcount;
+    *optional = argcount - min_argcount;
+
+    return OK;
+}
+
 /*
  * Find a function by name, including "<lambda>123".
  * Check for "profile" and "debug" arguments and set"compile_type".

src/version.c

@@ -704,6 +704,8 @@ static char *(features[]) =
 
 static int included_patches[] =
 {   /* Add new patch number below this line */
+/**/
+    547,
 /**/
     546,
 /**/