Add feature strip_directives (#217)

* Add feature `strip_directives`

This PR adds the `strip_directives` feature, a hotfix that will allow
yamlfmt to function with directives in the document for some cases. It's
a very imperfect feature but should help in some minimal cases.

* add missing license header

docs/config-file.md

@@ -77,17 +77,30 @@ The basic formatter is a barebones formatter that simply takes the data provided
 | `retain_line_breaks`        | bool           | false   | Retain line breaks in formatted yaml. |
 | `retain_line_breaks_single` | bool           | false   | (NOTE: Takes precedence over `retain_line_breaks`) Retain line breaks in formatted yaml, but only keep a single line in groups of many blank lines. |
 | `disallow_anchors`          | bool           | false   | If true, reject any YAML anchors or aliases found in the document. |
-| `max_line_length`           | int            | 0       | Set the maximum line length (see notes below). if not set, defaults to 0 which means no limit. |
+| `max_line_length`           | int            | 0       | Set the maximum line length ([see note below](#max_line_length)). if not set, defaults to 0 which means no limit. |
 | `scan_folded_as_literal`    | bool           | false   | Option that will preserve newlines in folded block scalars (blocks that start with `>`). |
 | `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   | 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. |
+| `strip_directives`          | bool           | false   | [YAML Directives](https://yaml.org/spec/1.2.2/#3234-directives) are not supported by this formatter. This feature will attempt to strip the directives before formatting and put them back. [Use this feature at your own risk.](#strip_directives) |
 
-### Note on `max_line_length`
+## Additional Notes
+
+### `max_line_length`
 
 It's not perfect; it uses the `best_width` setting from the [gopkg.in/yaml.v3][1] library. If there's a very long token that extends too far for the line width, it won't split it up properly. I will keep trying to make this work better, but decided to get a version of the feature in that works for a lot of scenarios even if not all of them.
 
+### `strip_directives`
+
+TL;DR:
+* If you only have directives at the top of the file this feature will work just fine, otherwise make sure you test it first.
+* Please note that directives are completely tossed and ignored by the formatter
+
+This hotfix is flaky. It is very hard to reconstruct data like this without parsing or knowing what may have changed about the structure of the document. What it attempts to do is find the directives in the document before formatting, keep track of them, and put them back where they "belong". However, the only mechanism it has to decide where it "belongs" is the line it was at originally. This can easily change based on what the formatter ended up changing. This means that the only way this fix predictably and reliably works is for directives that are at the top of the document before the document actually starts (i.e. where the `%YAML` directive goes).
+
+In addition, while with this feature the `%YAML` directive may work, the formatter very specifically supports only the [YAML 1.2 spec](https://yaml.org/spec/1.2.2/). So the `%YAML:1.0` directive won't have the desired effect when passing a file through `yamlfmt`, and if you have 1.0-only syntax in your document the formatter may end up failing in other ways that will be unfixable.
+
 [1]: https://www.github.com/braydonk/yaml
 [Specifying Paths]: ./paths.md

feature.go

@@ -14,9 +14,12 @@
 
 package yamlfmt
 
-import "fmt"
+import (
+	"context"
+	"fmt"
+)
 
-type FeatureFunc func([]byte) ([]byte, error)
+type FeatureFunc func(context.Context, []byte) (context.Context, []byte, error)
 
 type Feature struct {
 	Name         string
@@ -47,29 +50,29 @@ func (e *FeatureApplyError) Unwrap() error {
 	return e.err
 }
 
-func (fl FeatureList) ApplyFeatures(input []byte, mode FeatureApplyMode) ([]byte, error) {
+func (fl FeatureList) ApplyFeatures(ctx context.Context, input []byte, mode FeatureApplyMode) (context.Context, []byte, error) {
 	// Declare err here so the result variable doesn't get shadowed in the loop
 	var err error
 	result := make([]byte, len(input))
 	copy(result, input)
 	for _, feature := range fl {
 		if mode == FeatureApplyBefore {
 			if feature.BeforeAction != nil {
-				result, err = feature.BeforeAction(result)
+				ctx, result, err = feature.BeforeAction(ctx, result)
 			}
 		} else {
 			if feature.AfterAction != nil {
-				result, err = feature.AfterAction(result)
+				ctx, result, err = feature.AfterAction(ctx, result)
 			}
 		}
 
 		if err != nil {
-			return nil, &FeatureApplyError{
+			return ctx, nil, &FeatureApplyError{
 				err:         err,
 				featureName: feature.Name,
 				mode:        mode,
 			}
 		}
 	}
-	return result, nil
+	return ctx, result, nil
 }

formatters/basic/config.go

@@ -34,6 +34,7 @@ type Config struct {
 	PadLineComments        int                    `mapstructure:"pad_line_comments"`
 	TrimTrailingWhitespace bool                   `mapstructure:"trim_trailing_whitespace"`
 	EOFNewline             bool                   `mapstructure:"eof_newline"`
+	StripDirectives        bool                   `mapstructure:"strip_directives"`
 }
 
 func DefaultConfig() *Config {

formatters/basic/features.go

@@ -46,6 +46,12 @@ func ConfigureFeaturesFromConfig(config *Config) yamlfmt.FeatureList {
 			features.MakeFeatureEOFNewline(lineSep),
 		)
 	}
+	if config.StripDirectives {
+		configuredFeatures = append(
+			configuredFeatures,
+			hotfix.MakeFeatureStripDirectives(lineSep),
+		)
+	}
 	return configuredFeatures
 }
 

formatters/basic/formatter.go

@@ -16,6 +16,7 @@ package basic
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"io"
 
@@ -40,7 +41,8 @@ func (f *BasicFormatter) Type() string {
 
 func (f *BasicFormatter) Format(input []byte) ([]byte, error) {
 	// Run all features with BeforeActions
-	yamlContent, err := f.Features.ApplyFeatures(input, yamlfmt.FeatureApplyBefore)
+	ctx := context.Background()
+	ctx, yamlContent, err := f.Features.ApplyFeatures(ctx, input, yamlfmt.FeatureApplyBefore)
 	if err != nil {
 		return nil, err
 	}
@@ -78,7 +80,7 @@ func (f *BasicFormatter) Format(input []byte) ([]byte, error) {
 	}
 
 	// Run all features with AfterActions
-	resultYaml, err := f.Features.ApplyFeatures(b.Bytes(), yamlfmt.FeatureApplyAfter)
+	_, resultYaml, err := f.Features.ApplyFeatures(ctx, b.Bytes(), yamlfmt.FeatureApplyAfter)
 	if err != nil {
 		return nil, err
 	}

formatters/basic/formatter_test.go

@@ -20,6 +20,7 @@ import (
 
 	"github.com/google/yamlfmt"
 	"github.com/google/yamlfmt/formatters/basic"
+	"github.com/google/yamlfmt/internal/assert"
 )
 
 func newFormatter(config *basic.Config) *basic.BasicFormatter {
@@ -351,3 +352,14 @@ b: 2
 		t.Fatalf("expected: '%s', got: '%s'", expectedYml, resultStr)
 	}
 }
+
+func TestStripDirectives(t *testing.T) {
+	config := basic.DefaultConfig()
+	config.StripDirectives = true
+	f := newFormatter(config)
+
+	yml := "%YAML:1.0"
+
+	_, err := f.Format([]byte(yml))
+	assert.NilErr(t, err)
+}

integrationtest/command/command_test.go

@@ -139,3 +139,11 @@ func TestEOFNewline(t *testing.T) {
 		Update:  *updateFlag,
 	}.Run(t)
 }
+
+func TestStripDirectives(t *testing.T) {
+	TestCase{
+		Dir:     "strip_directives",
+		Command: yamlfmtWithArgs("-formatter strip_directives=true ."),
+		Update:  *updateFlag,
+	}.Run(t)
+}

integrationtest/command/testdata/print_conf_file/stdout/stdout.txt

@@ -24,5 +24,6 @@ formatter:
     retain_line_breaks: false
     retain_line_breaks_single: true
     scan_folded_as_literal: false
+    strip_directives: false
     trim_trailing_whitespace: false
     type: basic

integrationtest/command/testdata/print_conf_flags/stdout/stdout.txt

@@ -23,5 +23,6 @@ formatter:
     retain_line_breaks: true
     retain_line_breaks_single: false
     scan_folded_as_literal: false
+    strip_directives: false
     trim_trailing_whitespace: false
     type: basic

integrationtest/command/testdata/print_conf_flags_and_file/stdout/stdout.txt

@@ -24,5 +24,6 @@ formatter:
     retain_line_breaks: true
     retain_line_breaks_single: true
     scan_folded_as_literal: false
+    strip_directives: false
     trim_trailing_whitespace: false
     type: basic

integrationtest/command/testdata/strip_directives/after/a.yaml

@@ -0,0 +1,4 @@
+%YAML:1.0
+a: 1
+%TAG:hi
+b: 2

integrationtest/command/testdata/strip_directives/before/a.yaml

@@ -0,0 +1,4 @@
+%YAML:1.0
+a: 1
+%TAG:hi
+b: 2

internal/features/eof_newline.go

@@ -15,6 +15,8 @@
 package features
 
 import (
+	"context"
+
 	"github.com/google/yamlfmt"
 )
 
@@ -26,12 +28,12 @@ func MakeFeatureEOFNewline(linebreakStr string) yamlfmt.Feature {
 }
 
 func eofNewlineFeature(linebreakStr string) yamlfmt.FeatureFunc {
-	return func(content []byte) ([]byte, error) {
+	return func(_ context.Context, content []byte) (context.Context, []byte, error) {
 		// This check works in both linebreak modes.
 		if len(content) == 0 || content[len(content)-1] != '\n' {
 			linebreakBytes := []byte(linebreakStr)
 			content = append(content, linebreakBytes...)
 		}
-		return content, nil
+		return nil, content, nil
 	}
 }

internal/features/trim_whitespace.go

@@ -17,6 +17,7 @@ package features
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"strings"
 
 	"github.com/google/yamlfmt"
@@ -30,13 +31,13 @@ func MakeFeatureTrimTrailingWhitespace(linebreakStr string) yamlfmt.Feature {
 }
 
 func trimTrailingWhitespaceFeature(linebreakStr string) yamlfmt.FeatureFunc {
-	return func(content []byte) ([]byte, error) {
+	return func(_ context.Context, content []byte) (context.Context, []byte, error) {
 		buf := bytes.NewBuffer(content)
 		s := bufio.NewScanner(buf)
 		newLines := []string{}
 		for s.Scan() {
 			newLines = append(newLines, strings.TrimRight(s.Text(), " "))
 		}
-		return []byte(strings.Join(newLines, linebreakStr)), nil
+		return nil, []byte(strings.Join(newLines, linebreakStr)), nil
 	}
 }

internal/hotfix/retain_line_break.go

@@ -20,6 +20,7 @@ package hotfix
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"strings"
 
 	"github.com/google/yamlfmt"
@@ -51,7 +52,7 @@ func MakeFeatureRetainLineBreak(linebreakStr string, chomp bool) yamlfmt.Feature
 }
 
 func replaceLineBreakFeature(newlineStr string, chomp bool) yamlfmt.FeatureFunc {
-	return func(content []byte) ([]byte, error) {
+	return func(_ context.Context, content []byte) (context.Context, []byte, error) {
 		var buf bytes.Buffer
 		reader := bytes.NewReader(content)
 		scanner := bufio.NewScanner(reader)
@@ -74,12 +75,12 @@ func replaceLineBreakFeature(newlineStr string, chomp bool) yamlfmt.FeatureFunc
 				inLineBreaks = false
 			}
 		}
-		return buf.Bytes(), scanner.Err()
+		return nil, buf.Bytes(), scanner.Err()
 	}
 }
 
 func restoreLineBreakFeature(newlineStr string) yamlfmt.FeatureFunc {
-	return func(content []byte) ([]byte, error) {
+	return func(_ context.Context, content []byte) (context.Context, []byte, error) {
 		var buf bytes.Buffer
 		reader := bytes.NewReader(content)
 		scanner := bufio.NewScanner(reader)
@@ -98,6 +99,6 @@ func restoreLineBreakFeature(newlineStr string) yamlfmt.FeatureFunc {
 			buf.WriteString(txt)
 			buf.WriteString(newlineStr)
 		}
-		return buf.Bytes(), scanner.Err()
+		return nil, buf.Bytes(), scanner.Err()
 	}
 }

internal/hotfix/strip_directives.go

@@ -0,0 +1,101 @@
+// 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 hotfix
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"strings"
+
+	"github.com/google/yamlfmt"
+)
+
+type directiveKey string
+
+var contextDirectivesKey directiveKey = "directives"
+
+type Directive struct {
+	line    int
+	content string
+}
+
+func ContextWithDirectives(ctx context.Context, directives []Directive) context.Context {
+	return context.WithValue(ctx, contextDirectivesKey, directives)
+}
+
+func DirectivesFromContext(ctx context.Context) []Directive {
+	return ctx.Value(contextDirectivesKey).([]Directive)
+}
+
+func MakeFeatureStripDirectives(lineSepChar string) yamlfmt.Feature {
+	return yamlfmt.Feature{
+		Name:         "Strip Directives",
+		BeforeAction: stripDirectivesFeature(lineSepChar),
+		AfterAction:  restoreDirectivesFeature(lineSepChar),
+	}
+}
+
+func stripDirectivesFeature(lineSepChar string) yamlfmt.FeatureFunc {
+	return func(ctx context.Context, content []byte) (context.Context, []byte, error) {
+		directives := []Directive{}
+		reader := bytes.NewReader(content)
+		scanner := bufio.NewScanner(reader)
+		result := ""
+		currLine := 1
+		for scanner.Scan() {
+			line := scanner.Text()
+			if strings.HasPrefix(line, "%") {
+				directives = append(directives, Directive{
+					line:    currLine,
+					content: line,
+				})
+			} else {
+				result += line + lineSepChar
+			}
+			currLine++
+		}
+		return ContextWithDirectives(ctx, directives), []byte(result), nil
+	}
+}
+
+func restoreDirectivesFeature(lineSepChar string) yamlfmt.FeatureFunc {
+	return func(ctx context.Context, content []byte) (context.Context, []byte, error) {
+		directives := DirectivesFromContext(ctx)
+		directiveIdx := 0
+		doneDirectives := directiveIdx == len(directives)
+		reader := bytes.NewReader(content)
+		scanner := bufio.NewScanner(reader)
+		result := ""
+		currLine := 1
+		for scanner.Scan() {
+			if !doneDirectives && currLine == directives[directiveIdx].line {
+				result += directives[directiveIdx].content + lineSepChar
+				currLine++
+				directiveIdx++
+				doneDirectives = directiveIdx == len(directives)
+			}
+			result += scanner.Text() + lineSepChar
+			currLine++
+		}
+		// Edge case: There technically can be a directive as the final line. This would be
+		// useless as far as I can tell so maybe yamlfmt should just remove it anyway LOL but
+		// no we'll keep it.
+		if !doneDirectives && currLine == directives[directiveIdx].line {
+			result += directives[directiveIdx].content + lineSepChar
+		}
+		return ctx, []byte(result), nil
+	}
+}