From 7202c121a594cf2aabfce1c5ef6d0ced22664c73 Mon Sep 17 00:00:00 2001
From: Braydon Kains <93549768+braydonk@users.noreply.github.com>
Date: Mon, 24 Jun 2024 11:49:01 -0400
Subject: [PATCH] Add eof_newline option (#189)

This PR adds a new option to forcibly add the end of file newline if
it's not there. This is useful in the scenario where the
`retain_line_breaks` feature is disabled but the eof newline is still
wanted.
---
 docs/config-file.md                           |  3 +-
 formatters/basic/config.go                    |  1 +
 formatters/basic/features.go                  | 22 +++++++----
 formatters/basic/formatter_test.go            | 22 +++++++++++
 integrationtest/command/command_test.go       |  8 ++++
 .../command/testdata/eof_newline/after/a.yaml |  2 +
 .../testdata/eof_newline/before/a.yaml        |  2 +
 .../testdata/eof_newline/stdout/stderr.txt    |  0
 .../testdata/eof_newline/stdout/stdout.txt    |  0
 .../print_conf_file/stdout/stdout.txt         |  1 +
 .../print_conf_flags/stdout/stdout.txt        |  1 +
 .../stdout/stdout.txt                         |  1 +
 internal/features/eof_newline.go              | 37 +++++++++++++++++++
 .../{trim => features}/trim_whitespace.go     |  2 +-
 14 files changed, 92 insertions(+), 10 deletions(-)
 create mode 100755 integrationtest/command/testdata/eof_newline/after/a.yaml
 create mode 100644 integrationtest/command/testdata/eof_newline/before/a.yaml
 create mode 100755 integrationtest/command/testdata/eof_newline/stdout/stderr.txt
 create mode 100755 integrationtest/command/testdata/eof_newline/stdout/stdout.txt
 create mode 100644 internal/features/eof_newline.go
 rename internal/{trim => features}/trim_whitespace.go (98%)

diff --git a/docs/config-file.md b/docs/config-file.md
index 54b6d19..6a86260 100644
--- a/docs/config-file.md
+++ b/docs/config-file.md
@@ -82,7 +82,8 @@ The basic formatter is a barebones formatter that simply takes the data provided
 | `indentless_arrays`         | bool           | false   | Render `-` array items (block sequence items) without an increased indent. |
 | `drop_merge_tag`            | bool           | false   | Assume that any well formed merge using just a `<<` token will be a merge, and drop the `!!merge` tag from the formatted result. |
 | `pad_line_comments`         | int            | 1       | The number of padding spaces to insert before line comments. |
-| `trim_trailing_whitespace`  | bool           | false   | Whether to trim trailing whitespace from lines. |
+| `trim_trailing_whitespace`  | bool           | false   | Trim trailing whitespace from lines. |
+| `eof_newline`               | bool           | false   | Always add a newline at end of file. Useful in the scenario where `retain_line_breaks` is disabled but the trailing newline is still needed. |
 
 ### Note on `max_line_length`
 
diff --git a/formatters/basic/config.go b/formatters/basic/config.go
index a9e4959..b51a49d 100644
--- a/formatters/basic/config.go
+++ b/formatters/basic/config.go
@@ -33,6 +33,7 @@ type Config struct {
 	DropMergeTag           bool                   `mapstructure:"drop_merge_tag"`
 	PadLineComments        int                    `mapstructure:"pad_line_comments"`
 	TrimTrailingWhitespace bool                   `mapstructure:"trim_trailing_whitespace"`
+	EOFNewline             bool                   `mapstructure:"eof_newline"`
 }
 
 func DefaultConfig() *Config {
diff --git a/formatters/basic/features.go b/formatters/basic/features.go
index 0f0f382..a638191 100644
--- a/formatters/basic/features.go
+++ b/formatters/basic/features.go
@@ -18,8 +18,8 @@ import (
 	"github.com/braydonk/yaml"
 	"github.com/google/yamlfmt"
 	"github.com/google/yamlfmt/formatters/basic/anchors"
+	"github.com/google/yamlfmt/internal/features"
 	"github.com/google/yamlfmt/internal/hotfix"
-	"github.com/google/yamlfmt/internal/trim"
 )
 
 func ConfigureFeaturesFromConfig(config *Config) yamlfmt.FeatureList {
@@ -27,20 +27,26 @@ func ConfigureFeaturesFromConfig(config *Config) yamlfmt.FeatureList {
 	if err != nil {
 		lineSep = "\n"
 	}
-	features := []yamlfmt.Feature{}
+	configuredFeatures := []yamlfmt.Feature{}
 	if config.RetainLineBreaks || config.RetainLineBreaksSingle {
-		features = append(
-			features,
+		configuredFeatures = append(
+			configuredFeatures,
 			hotfix.MakeFeatureRetainLineBreak(lineSep, config.RetainLineBreaksSingle),
 		)
 	}
 	if config.TrimTrailingWhitespace {
-		features = append(
-			features,
-			trim.MakeFeatureTrimTrailingWhitespace(lineSep),
+		configuredFeatures = append(
+			configuredFeatures,
+			features.MakeFeatureTrimTrailingWhitespace(lineSep),
 		)
 	}
-	return features
+	if config.EOFNewline {
+		configuredFeatures = append(
+			configuredFeatures,
+			features.MakeFeatureEOFNewline(lineSep),
+		)
+	}
+	return configuredFeatures
 }
 
 // These features will directly use the `yaml.Node` type and
diff --git a/formatters/basic/formatter_test.go b/formatters/basic/formatter_test.go
index 7ae9ba1..ba01175 100644
--- a/formatters/basic/formatter_test.go
+++ b/formatters/basic/formatter_test.go
@@ -329,3 +329,25 @@ b: 2`
 		t.Fatalf("expected: '%s', got: '%s'", expectedYml, resultStr)
 	}
 }
+
+func TestEOFNewline(t *testing.T) {
+	config := basic.DefaultConfig()
+	config.RetainLineBreaks = false
+	config.EOFNewline = true
+	f := newFormatter(config)
+
+	yml := `a: 1
+b: 2`
+	expectedYml := `a: 1
+b: 2
+`
+
+	result, err := f.Format([]byte(yml))
+	if err != nil {
+		t.Fatalf("expected formatting to pass, returned error: %v", err)
+	}
+	resultStr := string(result)
+	if resultStr != expectedYml {
+		t.Fatalf("expected: '%s', got: '%s'", expectedYml, resultStr)
+	}
+}
diff --git a/integrationtest/command/command_test.go b/integrationtest/command/command_test.go
index 2b899e4..bd80839 100644
--- a/integrationtest/command/command_test.go
+++ b/integrationtest/command/command_test.go
@@ -131,3 +131,11 @@ func TestMultilineStringBug(t *testing.T) {
 		Update:  *updateFlag,
 	}.Run(t)
 }
+
+func TestEOFNewline(t *testing.T) {
+	TestCase{
+		Dir:     "eof_newline",
+		Command: yamlfmtWithArgs("-formatter eof_newline=true ."),
+		Update:  *updateFlag,
+	}.Run(t)
+}
diff --git a/integrationtest/command/testdata/eof_newline/after/a.yaml b/integrationtest/command/testdata/eof_newline/after/a.yaml
new file mode 100755
index 0000000..ce46e30
--- /dev/null
+++ b/integrationtest/command/testdata/eof_newline/after/a.yaml
@@ -0,0 +1,2 @@
+# no newline at end
+a: 1
diff --git a/integrationtest/command/testdata/eof_newline/before/a.yaml b/integrationtest/command/testdata/eof_newline/before/a.yaml
new file mode 100644
index 0000000..e4c3649
--- /dev/null
+++ b/integrationtest/command/testdata/eof_newline/before/a.yaml
@@ -0,0 +1,2 @@
+# no newline at end
+a: 1
\ No newline at end of file
diff --git a/integrationtest/command/testdata/eof_newline/stdout/stderr.txt b/integrationtest/command/testdata/eof_newline/stdout/stderr.txt
new file mode 100755
index 0000000..e69de29
diff --git a/integrationtest/command/testdata/eof_newline/stdout/stdout.txt b/integrationtest/command/testdata/eof_newline/stdout/stdout.txt
new file mode 100755
index 0000000..e69de29
diff --git a/integrationtest/command/testdata/print_conf_file/stdout/stdout.txt b/integrationtest/command/testdata/print_conf_file/stdout/stdout.txt
index d8598e4..af6692d 100644
--- a/integrationtest/command/testdata/print_conf_file/stdout/stdout.txt
+++ b/integrationtest/command/testdata/print_conf_file/stdout/stdout.txt
@@ -14,6 +14,7 @@ regex_exclude: []
 formatter:
     disallow_anchors: false
     drop_merge_tag: false
+    eof_newline: false
     include_document_start: true
     indent: 2
     indentless_arrays: false
diff --git a/integrationtest/command/testdata/print_conf_flags/stdout/stdout.txt b/integrationtest/command/testdata/print_conf_flags/stdout/stdout.txt
index fff8566..80ace34 100644
--- a/integrationtest/command/testdata/print_conf_flags/stdout/stdout.txt
+++ b/integrationtest/command/testdata/print_conf_flags/stdout/stdout.txt
@@ -13,6 +13,7 @@ regex_exclude: []
 formatter:
     disallow_anchors: false
     drop_merge_tag: false
+    eof_newline: false
     include_document_start: false
     indent: 2
     indentless_arrays: false
diff --git a/integrationtest/command/testdata/print_conf_flags_and_file/stdout/stdout.txt b/integrationtest/command/testdata/print_conf_flags_and_file/stdout/stdout.txt
index ce14d09..b6dd917 100644
--- a/integrationtest/command/testdata/print_conf_flags_and_file/stdout/stdout.txt
+++ b/integrationtest/command/testdata/print_conf_flags_and_file/stdout/stdout.txt
@@ -14,6 +14,7 @@ regex_exclude: []
 formatter:
     disallow_anchors: false
     drop_merge_tag: false
+    eof_newline: false
     include_document_start: true
     indent: 2
     indentless_arrays: false
diff --git a/internal/features/eof_newline.go b/internal/features/eof_newline.go
new file mode 100644
index 0000000..568db9b
--- /dev/null
+++ b/internal/features/eof_newline.go
@@ -0,0 +1,37 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package features
+
+import (
+	"github.com/google/yamlfmt"
+)
+
+func MakeFeatureEOFNewline(linebreakStr string) yamlfmt.Feature {
+	return yamlfmt.Feature{
+		Name:        "EOF Newline",
+		AfterAction: eofNewlineFeature(linebreakStr),
+	}
+}
+
+func eofNewlineFeature(linebreakStr string) yamlfmt.FeatureFunc {
+	return func(content []byte) ([]byte, error) {
+		// This check works in both linebreak modes.
+		if content[len(content)-1] != '\n' {
+			linebreakBytes := []byte(linebreakStr)
+			content = append(content, linebreakBytes...)
+		}
+		return content, nil
+	}
+}
diff --git a/internal/trim/trim_whitespace.go b/internal/features/trim_whitespace.go
similarity index 98%
rename from internal/trim/trim_whitespace.go
rename to internal/features/trim_whitespace.go
index 1bcbcf9..716faad 100644
--- a/internal/trim/trim_whitespace.go
+++ b/internal/features/trim_whitespace.go
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package trim
+package features
 
 import (
 	"bufio"