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.

Author: brammool

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.
+		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.
+		unmodified, also when binary mode is not specified.
 
-		When {flags} contains "a" then append mode is used, lines are
-		appended to the file: >
+		{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")
 <
-		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.
+		'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.
 
-		All NL characters are replaced with a NUL character.
-		Inserting CR characters needs to be done before passing {list}
-		to writefile().
 		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,37 +2300,59 @@ 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)
-    {
-	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;
+	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;
+	    }
+	    else
+	    {
+		if (write_list(fd, list, binary) == 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));
+	    if (ret == 0 && 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);
+	    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,33 +5571,58 @@ 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)
     {
 	semsg(_(e_str_not_inside_function), "defer");
 	return FAIL;
     }
     if (get_func_arguments(arg, evalarg, FALSE, argvars, &argcount) == FAIL)
-	goto theend;
-    saved_name = vim_strsave(name);
-    if (saved_name == NULL)
-	goto theend;
+    {
+	while (--argcount >= 0)
+	    clear_tv(&argvars[argcount]);
+	return FAIL;
+    }
+    return add_defer(name, argcount, argvars);
+}
 
-    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)
+/*
+ * 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;
-    dr = ((defer_T *)current_funccal->fc_defer.ga_data)
-					  + current_funccal->fc_defer.ga_len++;
-    dr->dr_name = saved_name;
-    dr->dr_argcount = argcount;
-    while (argcount > 0)
+    if (in_def_function())
+    {
+	if (add_defer_function(saved_name, argcount, argvars) == OK)
+	    argcount = 0;
+    }
+    else
     {
-	--argcount;
-	dr->dr_argvars[argcount] = argvars[argcount];
+	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)
+	    goto theend;
+	dr = ((defer_T *)current_funccal->fc_defer.ga_data)
+					  + current_funccal->fc_defer.ga_len++;
+	dr->dr_name = saved_name;
+	dr->dr_argcount = argcount;
+	while (argcount > 0)
+	{
+	    --argcount;
+	    dr->dr_argvars[argcount] = argvars[argcount];
+	}
     }
     ret = OK;
 

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;
-    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 NULL;
-	dfunc->df_defer_var_idx = lvar->lv_idx + 1;
-    }
-    if (generate_DEFER(cctx, dfunc->df_defer_var_idx - 1, argcount) == FAIL)
+    defer_var_idx = get_defer_var_idx(cctx);
+    if (defer_var_idx == 0)
+	return NULL;
+    if (generate_DEFER(cctx, defer_var_idx - 1, argcount) == FAIL)
 	return NULL;
 
     return skipwhite(arg);