From 745b6c0e5f10892d51275b204e2942fdf01c1202 Mon Sep 17 00:00:00 2001
From: Elara6331 <elara@elara.ws>
Date: Thu, 28 Dec 2023 07:19:58 -0800
Subject: [PATCH] Add WriteOnSuccess option

---
 namespace.go | 14 +++++++++++++-
 parse.go     | 11 ++++++-----
 salix.go     | 26 +++++++++++++++++++++++---
 3 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/namespace.go b/namespace.go
index e40e21a..a26b255 100644
--- a/namespace.go
+++ b/namespace.go
@@ -17,7 +17,10 @@ type Namespace struct {
 	// to make the resulting document look better. Postprocessing is only done once when the
 	// template is parsed, so it will not affect performance. (default: true)
 	WhitespaceMutations bool
-	escapeHTML          *bool
+	// WriteOnSuccess indicates whether the output should only be written if generation fully succeeds.
+	// This option buffers the output of the template, so it will use more memory. (default: false)
+	WriteOnSuccess bool
+	escapeHTML     *bool
 }
 
 // New returns a new template namespace
@@ -27,6 +30,7 @@ func New() *Namespace {
 		vars:                map[string]any{},
 		tags:                map[string]Tag{},
 		WhitespaceMutations: true,
+		WriteOnSuccess:      false,
 	}
 }
 
@@ -66,6 +70,14 @@ func (n *Namespace) WithEscapeHTML(b bool) *Namespace {
 	return n
 }
 
+// WithWriteOnSuccess enables or disables only writing if generation fully succeeds.
+func (n *Namespace) WithWriteOnSuccess(b bool) *Namespace {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.WriteOnSuccess = true
+	return n
+}
+
 // WithWhitespaceMutations turns whitespace mutations on or off for the namespace
 func (n *Namespace) WithWhitespaceMutations(b bool) *Namespace {
 	n.mu.Lock()
diff --git a/parse.go b/parse.go
index bc9484f..6042bb3 100644
--- a/parse.go
+++ b/parse.go
@@ -32,11 +32,12 @@ func (n *Namespace) ParseWithName(name string, r io.Reader) (Template, error) {
 	}
 
 	t := Template{
-		ns:   n,
-		name: name,
-		ast:  astVal.([]ast.Node),
-		tags: map[string]Tag{},
-		vars: map[string]any{},
+		ns:             n,
+		name:           name,
+		ast:            astVal.([]ast.Node),
+		tags:           map[string]Tag{},
+		vars:           map[string]any{},
+		WriteOnSuccess: n.WriteOnSuccess,
 	}
 
 	if n.WhitespaceMutations {
diff --git a/salix.go b/salix.go
index 0ac086c..809037e 100644
--- a/salix.go
+++ b/salix.go
@@ -2,6 +2,7 @@ package salix
 
 import (
 	"bufio"
+	"bytes"
 	"errors"
 	"fmt"
 	"html"
@@ -22,6 +23,9 @@ type Template struct {
 	ast  []ast.Node
 
 	escapeHTML *bool
+	// WriteOnSuccess indicates whether the output should only be written if generation fully succeeds.
+	// This option buffers the output of the template, so it will use more memory. (default: false)
+	WriteOnSuccess bool
 
 	tags   map[string]Tag
 	vars   map[string]any
@@ -56,13 +60,29 @@ func (t Template) WithEscapeHTML(b bool) Template {
 	return t
 }
 
+// WithWriteOnSuccess enables or disables only writing if generation fully succeeds.
+func (t Template) WithWriteOnSuccess(b bool) Template {
+	t.WriteOnSuccess = true
+	return t
+}
+
 // Execute executes a parsed template and writes
 // the result to w.
 func (t Template) Execute(w io.Writer) error {
 	t.macros = map[string][]ast.Node{}
-	bw := bufio.NewWriterSize(w, 16384)
-	defer bw.Flush()
-	return t.execute(bw, t.ast, nil)
+	if t.WriteOnSuccess {
+		buf := &bytes.Buffer{}
+		err := t.execute(buf, t.ast, nil)
+		if err != nil {
+			return err
+		}
+		_, err = io.Copy(w, buf)
+		return err
+	} else {
+		bw := bufio.NewWriterSize(w, 16384)
+		defer bw.Flush()
+		return t.execute(bw, t.ast, nil)
+	}
 }
 
 func (t *Template) execute(w io.Writer, nodes []ast.Node, local map[string]any) error {