patch 9.0.0379: cleaning up after writefile() is a hassle

Problem:    Cleaning up after writefile() is a hassle.
Solution:   Add the 'D' flag to defer deleting the written file.  Very useful
            in tests.

runtime/doc/builtin.txt

@@ -10448,34 +10448,43 @@ writefile({object}, {fname} [, {flags}])
 		When {object} is a |List| write it to file {fname}.  Each list
 		item is separated with a NL.  Each list item must be a String
 		or Number.
-		When {flags} contains "b" then binary mode is used: There will
-		not be a NL after the last list item.  An empty item at the
-		end does cause the last line in the file to end in a NL.
-
-		When {object} is a |Blob| write the bytes to file {fname}
-		unmodified.
-
-		When {flags} contains "a" then append mode is used, lines are
-		appended to the file: >
-			:call writefile(["foo"], "event.log", "a")
-			:call writefile(["bar"], "event.log", "a")
-<
-		When {flags} contains "s" then fsync() is called after writing
-		the file.  This flushes the file to disk, if possible.  This
-		takes more time but avoids losing the file if the system
-		crashes.
-		When {flags} does not contain "S" or "s" then fsync() is
-		called if the 'fsync' option is set.
-		When {flags} contains "S" then fsync() is not called, even
-		when 'fsync' is set.
-
 		All NL characters are replaced with a NUL character.
 		Inserting CR characters needs to be done before passing {list}
 		to writefile().
+
+		When {object} is a |Blob| write the bytes to file {fname}
+		unmodified, also when binary mode is not specified.
+
+		{flags} must be a String.  These characters are recognized:
+
+		'b'  Binary mode is used: There will not be a NL after the
+		     last list item.  An empty item at the end does cause the
+		     last line in the file to end in a NL.
+
+		'a'  Append mode is used, lines are appended to the file: >
+			:call writefile(["foo"], "event.log", "a")
+			:call writefile(["bar"], "event.log", "a")
+<
+		'D'  Delete the file when the current function ends.  This
+		     works like: >
+		     	:defer delete({fname})
+<		     Fails when not in a function.  Also see |:defer|.
+
+		's'  fsync() is called after writing the file.  This flushes
+		     the file to disk, if possible.  This takes more time but
+		     avoids losing the file if the system crashes.
+
+		'S'  fsync() is not called, even when 'fsync' is set.
+
+		     When {flags} does not contain "S" or "s" then fsync() is
+		     called if the 'fsync' option is set.
+
 		An existing file is overwritten, if possible.
+
 		When the write fails -1 is returned, otherwise 0.  There is an
 		error message if the file can't be created or when writing
 		fails.
+
 		Also see |readfile()|.
 		To copy a file byte for byte: >
 			:let fl = readfile("foo", "b")

src/filepath.c

@@ -2232,6 +2232,7 @@ f_writefile(typval_T *argvars, typval_T *rettv)
 {
     int		binary = FALSE;
     int		append = FALSE;
+    int		defer = FALSE;
 #ifdef HAVE_FSYNC
     int		do_fsync = p_fs;
 #endif
@@ -2285,6 +2286,8 @@ f_writefile(typval_T *argvars, typval_T *rettv)
 	    binary = TRUE;
 	if (vim_strchr(arg2, 'a') != NULL)
 	    append = TRUE;
+	if (vim_strchr(arg2, 'D') != NULL)
+	    defer = TRUE;
 #ifdef HAVE_FSYNC
 	if (vim_strchr(arg2, 's') != NULL)
 	    do_fsync = TRUE;
@@ -2297,38 +2300,60 @@ f_writefile(typval_T *argvars, typval_T *rettv)
     if (fname == NULL)
 	return;
 
+    if (defer && !in_def_function() && get_current_funccal() == NULL)
+    {
+	semsg(_(e_str_not_inside_function), "defer");
+	return;
+    }
+
     // Always open the file in binary mode, library functions have a mind of
     // their own about CR-LF conversion.
     if (*fname == NUL || (fd = mch_fopen((char *)fname,
 				      append ? APPENDBIN : WRITEBIN)) == NULL)
     {
-	semsg(_(e_cant_create_file_str), *fname == NUL ? (char_u *)_("<empty>") : fname);
+	semsg(_(e_cant_create_file_str),
+			       *fname == NUL ? (char_u *)_("<empty>") : fname);
 	ret = -1;
     }
-    else if (blob)
+    else
+    {
+	if (defer)
+	{
+	    typval_T tv;
+
+	    tv.v_type = VAR_STRING;
+	    tv.v_lock = 0;
+	    tv.vval.v_string = vim_strsave(fname);
+	    if (tv.vval.v_string == NULL
+		    || add_defer((char_u *)"delete", 1, &tv) == FAIL)
+	    {
+		ret = -1;
+		fclose(fd);
+		(void)mch_remove(fname);
+	    }
+	}
+
+	if (ret == 0)
+	{
+	    if (blob)
 	    {
 		if (write_blob(fd, blob) == FAIL)
 		    ret = -1;
-#ifdef HAVE_FSYNC
-	else if (do_fsync)
-	    // Ignore the error, the user wouldn't know what to do about it.
-	    // May happen for a device.
-	    vim_ignored = vim_fsync(fileno(fd));
-#endif
-	fclose(fd);
 	    }
 	    else
 	    {
 		if (write_list(fd, list, binary) == FAIL)
 		    ret = -1;
+	    }
 #ifdef HAVE_FSYNC
-	else if (do_fsync)
+	    if (ret == 0 && do_fsync)
-	    // Ignore the error, the user wouldn't know what to do about it.
+		// Ignore the error, the user wouldn't know what to do about
-	    // May happen for a device.
+		// it.  May happen for a device.
 		vim_ignored = vim_fsync(fileno(fd));
 #endif
 	    fclose(fd);
 	}
+    }
 
     rettv->vval.v_number = ret;
 }

src/proto/userfunc.pro

@@ -58,6 +58,7 @@ void func_ptr_unref(ufunc_T *fp);
 void func_ref(char_u *name);
 void func_ptr_ref(ufunc_T *fp);
 void ex_return(exarg_T *eap);
+int add_defer(char_u *name, int argcount_arg, typval_T *argvars);
 void handle_defer(void);
 void ex_call(exarg_T *eap);
 int do_return(exarg_T *eap, int reanimate, int is_cmd, void *rettv);

src/proto/vim9cmds.pro

@@ -21,6 +21,7 @@ char_u *compile_finally(char_u *arg, cctx_T *cctx);
 char_u *compile_endtry(char_u *arg, cctx_T *cctx);
 char_u *compile_throw(char_u *arg, cctx_T *cctx);
 char_u *compile_eval(char_u *arg, cctx_T *cctx);
+int get_defer_var_idx(cctx_T *cctx);
 char_u *compile_defer(char_u *arg_start, cctx_T *cctx);
 char_u *compile_mult_expr(char_u *arg, int cmdidx, cctx_T *cctx);
 char_u *compile_put(char_u *arg, exarg_T *eap, cctx_T *cctx);

src/proto/vim9execute.pro

@@ -3,6 +3,8 @@ void to_string_error(vartype_T vartype);
 void update_has_breakpoint(ufunc_T *ufunc);
 void funcstack_check_refcount(funcstack_T *funcstack);
 int set_ref_in_funcstacks(int copyID);
+int in_def_function(void);
+int add_defer_function(char_u *name, int argcount, typval_T *argvars);
 char_u *char_from_string(char_u *str, varnumber_T index);
 char_u *string_slice(char_u *str, varnumber_T first, varnumber_T last, int exclusive);
 int fill_partial_and_closure(partial_T *pt, ufunc_T *ufunc, ectx_T *ectx);

src/testdir/test_writefile.vim

@@ -933,6 +933,23 @@ func Test_write_binary_file()
   call delete('Xwbfile3')
 endfunc
 
+func DoWriteDefer()
+  call writefile(['some text'], 'XdeferDelete', 'D')
+  call assert_equal(['some text'], readfile('XdeferDelete'))
+endfunc
+
+def DefWriteDefer()
+  writefile(['some text'], 'XdefdeferDelete', 'D')
+  assert_equal(['some text'], readfile('XdefdeferDelete'))
+enddef
+
+func Test_write_with_deferred_delete()
+  call DoWriteDefer()
+  call assert_equal('', glob('XdeferDelete'))
+  call DefWriteDefer()
+  call assert_equal('', glob('XdefdeferDelete'))
+endfunc
+
 " Check that buffer is written before triggering QuitPre
 func Test_wq_quitpre_autocommand()
   edit Xsomefile

src/userfunc.c

@@ -1728,6 +1728,7 @@ emsg_funcname(char *ermsg, char_u *name)
 /*
  * Get function arguments at "*arg" and advance it.
  * Return them in "*argvars[MAX_FUNC_ARGS + 1]" and the count in "argcount".
+ * On failure FAIL is returned but the "argvars[argcount]" are still set.
  */
     static int
 get_func_arguments(
@@ -5570,9 +5571,6 @@ ex_defer_inner(char_u *name, char_u **arg, evalarg_T *evalarg)
 {
     typval_T	argvars[MAX_FUNC_ARGS + 1];	// vars for arguments
     int		argcount = 0;			// number of arguments found
-    defer_T	*dr;
-    int		ret = FAIL;
-    char_u	*saved_name;
 
     if (current_funccal == NULL)
     {
@@ -5580,11 +5578,38 @@ ex_defer_inner(char_u *name, char_u **arg, evalarg_T *evalarg)
 	return FAIL;
     }
     if (get_func_arguments(arg, evalarg, FALSE, argvars, &argcount) == FAIL)
-	goto theend;
+    {
-    saved_name = vim_strsave(name);
+	while (--argcount >= 0)
+	    clear_tv(&argvars[argcount]);
+	return FAIL;
+    }
+    return add_defer(name, argcount, argvars);
+}
+
+/*
+ * Add a deferred call for "name" with arguments "argvars[argcount]".
+ * Consumes "argvars[]".
+ * Caller must check that in_def_function() returns TRUE or current_funccal is
+ * not NULL.
+ * Returns OK or FAIL.
+ */
+    int
+add_defer(char_u *name, int argcount_arg, typval_T *argvars)
+{
+    char_u	*saved_name = vim_strsave(name);
+    int		argcount = argcount_arg;
+    defer_T	*dr;
+    int		ret = FAIL;
+
     if (saved_name == NULL)
 	goto theend;
-
+    if (in_def_function())
+    {
+	if (add_defer_function(saved_name, argcount, argvars) == OK)
+	    argcount = 0;
+    }
+    else
+    {
 	if (current_funccal->fc_defer.ga_itemsize == 0)
 	    ga_init2(&current_funccal->fc_defer, sizeof(defer_T), 10);
 	if (ga_grow(&current_funccal->fc_defer, 1) == FAIL)
@@ -5598,6 +5623,7 @@ ex_defer_inner(char_u *name, char_u **arg, evalarg_T *evalarg)
 	    --argcount;
 	    dr->dr_argvars[argcount] = argvars[argcount];
 	}
+    }
     ret = OK;
 
 theend:

src/version.c

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

src/vim9cmds.c

@@ -1684,6 +1684,27 @@ compile_eval(char_u *arg, cctx_T *cctx)
     return skipwhite(p);
 }
 
+/*
+ * Get the local variable index for deferred function calls.
+ * Reserve it when not done already.
+ * Returns zero for failure.
+ */
+    int
+get_defer_var_idx(cctx_T *cctx)
+{
+    dfunc_T	*dfunc = ((dfunc_T *)def_functions.ga_data)
+					       + cctx->ctx_ufunc->uf_dfunc_idx;
+    if (dfunc->df_defer_var_idx == 0)
+    {
+	lvar_T *lvar = reserve_local(cctx, (char_u *)"@defer@", 7,
+							    TRUE, &t_list_any);
+	if (lvar == NULL)
+	    return 0;
+	dfunc->df_defer_var_idx = lvar->lv_idx + 1;
+    }
+    return dfunc->df_defer_var_idx;
+}
+
 /*
  * Compile "defer func(arg)".
  */
@@ -1693,7 +1714,7 @@ compile_defer(char_u *arg_start, cctx_T *cctx)
     char_u	*p;
     char_u	*arg = arg_start;
     int		argcount = 0;
-    dfunc_T	*dfunc;
+    int		defer_var_idx;
     type_T	*type;
     int		func_idx;
 
@@ -1730,16 +1751,10 @@ compile_defer(char_u *arg_start, cctx_T *cctx)
 
     // TODO: check argument count with "type"
 
-    dfunc = ((dfunc_T *)def_functions.ga_data) + cctx->ctx_ufunc->uf_dfunc_idx;
+    defer_var_idx = get_defer_var_idx(cctx);
-    if (dfunc->df_defer_var_idx == 0)
+    if (defer_var_idx == 0)
-    {
-	lvar_T *lvar = reserve_local(cctx, (char_u *)"@defer@", 7,
-							    TRUE, &t_list_any);
-	if (lvar == NULL)
 	return NULL;
-	dfunc->df_defer_var_idx = lvar->lv_idx + 1;
+    if (generate_DEFER(cctx, defer_var_idx - 1, argcount) == FAIL)
-    }
-    if (generate_DEFER(cctx, dfunc->df_defer_var_idx - 1, argcount) == FAIL)
 	return NULL;
 
     return skipwhite(arg);

src/vim9execute.c

@@ -845,41 +845,71 @@ set_ref_in_funcstacks(int copyID)
     return FALSE;
 }
 
+// Ugly static to avoid passing the execution context around through many
+// layers.
+static ectx_T *current_ectx = NULL;
+
 /*
- * Handle ISN_DEFER.  Stack has a function reference and "argcount" arguments.
+ * Return TRUE if currently executing a :def function.
- * The local variable that lists deferred functions is "var_idx".
+ * Can be used by builtin functions only.
- * Returns OK or FAIL.
  */
-    static int
+    int
-add_defer_func(int var_idx, int argcount, ectx_T *ectx)
+in_def_function(void)
+{
+    return current_ectx != NULL;
+}
+
+/*
+ * Add an entry for a deferred function call to the currently executing
+ * function.
+ * Return the list or NULL when failed.
+ */
+    static list_T *
+add_defer_item(int var_idx, int argcount, ectx_T *ectx)
 {
     typval_T	*defer_tv = STACK_TV_VAR(var_idx);
     list_T	*defer_l;
-    typval_T	*func_tv;
     list_T	*l;
-    int		i;
     typval_T	listval;
 
     if (defer_tv->v_type != VAR_LIST)
     {
 	// first time, allocate the list
 	if (rettv_list_alloc(defer_tv) == FAIL)
-	    return FAIL;
+	    return NULL;
     }
     defer_l = defer_tv->vval.v_list;
 
     l = list_alloc_with_items(argcount + 1);
     if (l == NULL)
-	return FAIL;
+	return NULL;
     listval.v_type = VAR_LIST;
     listval.vval.v_list = l;
     listval.v_lock = 0;
     if (list_insert_tv(defer_l, &listval, defer_l->lv_first) == FAIL)
     {
 	vim_free(l);
-	return FAIL;
+	return NULL;
     }
 
+    return l;
+}
+
+/*
+ * Handle ISN_DEFER.  Stack has a function reference and "argcount" arguments.
+ * The local variable that lists deferred functions is "var_idx".
+ * Returns OK or FAIL.
+ */
+    static int
+defer_command(int var_idx, int argcount, ectx_T *ectx)
+{
+    list_T	*l = add_defer_item(var_idx, argcount, ectx);
+    int		i;
+    typval_T	*func_tv;
+
+    if (l == NULL)
+	return FAIL;
+
     func_tv = STACK_TV_BOT(-argcount - 1);
     // TODO: check type is a funcref
     list_set_item(l, 0, func_tv);
@@ -890,6 +920,43 @@ add_defer_func(int var_idx, int argcount, ectx_T *ectx)
     return OK;
 }
 
+/*
+ * Add a deferred function "name" with one argument "arg_tv".
+ * Consumes "name", also on failure.
+ * Only to be called when in_def_function() returns TRUE.
+ */
+    int
+add_defer_function(char_u *name, int argcount, typval_T *argvars)
+{
+    dfunc_T	*dfunc = ((dfunc_T *)def_functions.ga_data)
+						  + current_ectx->ec_dfunc_idx;
+    list_T	*l;
+    typval_T	func_tv;
+    int		i;
+
+    if (dfunc->df_defer_var_idx == 0)
+    {
+	iemsg("df_defer_var_idx is zero");
+	vim_free(func_tv.vval.v_string);
+	return FAIL;
+    }
+    func_tv.v_type = VAR_FUNC;
+    func_tv.v_lock = 0;
+    func_tv.vval.v_string = name;
+
+    l = add_defer_item(dfunc->df_defer_var_idx - 1, 1, current_ectx);
+    if (l == NULL)
+    {
+	vim_free(func_tv.vval.v_string);
+	return FAIL;
+    }
+
+    list_set_item(l, 0, &func_tv);
+    for (i = 0; i < argcount; ++i)
+	list_set_item(l, i + 1, argvars + i);
+    return OK;
+}
+
 /*
  * Invoked when returning from a function: Invoke any deferred calls.
  */
@@ -1068,10 +1135,6 @@ call_prepare(int argcount, typval_T *argvars, ectx_T *ectx)
     return OK;
 }
 
-// Ugly global to avoid passing the execution context around through many
-// layers.
-static ectx_T *current_ectx = NULL;
-
 /*
  * Call a builtin function by index.
  */
@@ -3748,7 +3811,7 @@ exec_instructions(ectx_T *ectx)
 
 	    // :defer func(arg)
 	    case ISN_DEFER:
-		if (add_defer_func(iptr->isn_arg.defer.defer_var_idx,
+		if (defer_command(iptr->isn_arg.defer.defer_var_idx,
 			     iptr->isn_arg.defer.defer_argcount, ectx) == FAIL)
 		    goto on_error;
 		break;
@@ -5121,7 +5184,7 @@ theend:
 }
 
 /*
- * Execute the instructions from a VAR_INSTR typeval and put the result in
+ * Execute the instructions from a VAR_INSTR typval and put the result in
  * "rettv".
  * Return OK or FAIL.
  */

src/vim9expr.c

@@ -833,6 +833,14 @@ compile_call(
 		}
 	    }
 
+	    if (STRCMP(name, "writefile") == 0 && argcount > 2)
+	    {
+		// May have the "D" flag, reserve a variable for a deferred
+		// function call.
+		if (get_defer_var_idx(cctx) == 0)
+		    idx = -1;
+	    }
+
 	    if (idx >= 0)
 		res = generate_BCALL(cctx, idx, argcount, argcount_init == 1);
 	}