From bcae65770a33bdadd8851d58b6817c6fef0ac02d Mon Sep 17 00:00:00 2001
From: Katrina Verey <katrina.verey@shopify.com>
Date: Thu, 21 Apr 2022 18:28:13 -0400
Subject: [PATCH] Configurable extensions for template parser

---
 kyaml/fn/framework/parser/parser.go        | 21 +++++++++++++++------
 kyaml/fn/framework/parser/schema.go        |  2 +-
 kyaml/fn/framework/parser/schema_test.go   |  2 +-
 kyaml/fn/framework/parser/template.go      |  8 +++++++-
 kyaml/fn/framework/parser/template_test.go | 19 ++++++++++++++++++-
 5 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/kyaml/fn/framework/parser/parser.go b/kyaml/fn/framework/parser/parser.go
index 748964486..a5bc7729c 100644
--- a/kyaml/fn/framework/parser/parser.go
+++ b/kyaml/fn/framework/parser/parser.go
@@ -13,9 +13,9 @@ import (
 )
 
 type parser struct {
-	fs        iofs.FS
-	paths     []string
-	extension string
+	fs         iofs.FS
+	paths      []string
+	extensions []string
 }
 
 type contentProcessor func(content []byte, name string) error
@@ -51,8 +51,8 @@ func (l parser) readPath(path string, processContent contentProcessor) error {
 	}
 
 	// Path is a file -- check extension and read it
-	if !strings.HasSuffix(path, l.extension) {
-		return errors.Errorf("file %s did not have required extension %s", path, l.extension)
+	if err := checkExtension(path, l.extensions); err != nil {
+		return err
 	}
 	b, err := ioutil.ReadAll(f)
 	if err != nil {
@@ -61,6 +61,15 @@ func (l parser) readPath(path string, processContent contentProcessor) error {
 	return processContent(b, path)
 }
 
+func checkExtension(path string, extensions []string) error {
+	for _, ext := range extensions {
+		if strings.HasSuffix(path, ext) {
+			return nil
+		}
+	}
+	return errors.Errorf("file %s does not have any of permitted extensions %s", path, extensions)
+}
+
 func (l parser) readDir(dir iofs.ReadDirFile, dirname string, processContent contentProcessor) error {
 	entries, err := dir.ReadDir(0)
 	if err != nil {
@@ -68,7 +77,7 @@ func (l parser) readDir(dir iofs.ReadDirFile, dirname string, processContent con
 	}
 
 	for _, entry := range entries {
-		if entry.IsDir() || !strings.HasSuffix(entry.Name(), l.extension) {
+		if err := checkExtension(entry.Name(), l.extensions); err != nil || entry.IsDir() {
 			continue
 		}
 		// Note: using filepath.Join will break Windows, because io/fs.FS implementations require slashes on all OS.
diff --git a/kyaml/fn/framework/parser/schema.go b/kyaml/fn/framework/parser/schema.go
index 2160178a9..d2c12840b 100644
--- a/kyaml/fn/framework/parser/schema.go
+++ b/kyaml/fn/framework/parser/schema.go
@@ -56,7 +56,7 @@ func SchemaStrings(data ...string) framework.SchemaParser {
 //		AdditionalSchemas: parser.SchemaFiles("path/to/crd-schemas", "path/to/special-schema.json),
 //	}
 func SchemaFiles(paths ...string) SchemaParser {
-	return SchemaParser{parser{paths: paths, extension: SchemaExtension}}
+	return SchemaParser{parser{paths: paths, extensions: []string{SchemaExtension}}}
 }
 
 // SchemaParser is a framework.SchemaParser that can parse files or directories containing openapi schemas.
diff --git a/kyaml/fn/framework/parser/schema_test.go b/kyaml/fn/framework/parser/schema_test.go
index a51801a95..9d78cd4c7 100644
--- a/kyaml/fn/framework/parser/schema_test.go
+++ b/kyaml/fn/framework/parser/schema_test.go
@@ -52,7 +52,7 @@ func TestSchemaFiles(t *testing.T) {
 		{
 			name:    "rejects non-.template.yaml files",
 			paths:   []string{"testdata/ignore.yaml"},
-			wantErr: "file testdata/ignore.yaml did not have required extension .json",
+			wantErr: "file testdata/ignore.yaml does not have any of permitted extensions [.json]",
 		},
 	}
 	for _, tt := range tests {
diff --git a/kyaml/fn/framework/parser/template.go b/kyaml/fn/framework/parser/template.go
index 44a6db436..b98381bd3 100644
--- a/kyaml/fn/framework/parser/template.go
+++ b/kyaml/fn/framework/parser/template.go
@@ -61,7 +61,7 @@ func TemplateStrings(data ...string) framework.TemplateParser {
 //		}},
 //   }
 func TemplateFiles(paths ...string) TemplateParser {
-	return TemplateParser{parser{paths: paths, extension: TemplateExtension}}
+	return TemplateParser{parser{paths: paths, extensions: []string{TemplateExtension}}}
 }
 
 // TemplateParser is a framework.TemplateParser that can parse files or directories containing Go templated YAML.
@@ -95,3 +95,9 @@ func (l TemplateParser) FromFS(fs iofs.FS) TemplateParser {
 	l.parser.fs = fs
 	return l
 }
+
+// WithExtensions allows you to replace the extension the parser accept on the input files.
+func (l TemplateParser) WithExtensions(ext ...string) TemplateParser {
+	l.parser.extensions = ext
+	return l
+}
diff --git a/kyaml/fn/framework/parser/template_test.go b/kyaml/fn/framework/parser/template_test.go
index 1b0e14605..bb0d60f07 100644
--- a/kyaml/fn/framework/parser/template_test.go
+++ b/kyaml/fn/framework/parser/template_test.go
@@ -81,7 +81,7 @@ func TestTemplateFiles(t *testing.T) {
 		{
 			name:    "rejects non-.template.yaml files",
 			paths:   []string{"testdata/ignore.yaml"},
-			wantErr: "file testdata/ignore.yaml did not have required extension .template.yaml",
+			wantErr: "file testdata/ignore.yaml does not have any of permitted extensions [.template.yaml]",
 		},
 	}
 	for _, tt := range tests {
@@ -114,6 +114,23 @@ func TestTemplateFiles(t *testing.T) {
 	}
 }
 
+func TestTemplateParser_WithExtensions(t *testing.T) {
+	p := parser.TemplateFiles("testdata").WithExtensions(".json")
+	templates, err := p.Parse()
+	require.NoError(t, err)
+	require.Len(t, templates, 2)
+
+	p = p.WithExtensions(".yaml")
+	templates, err = p.Parse()
+	require.NoError(t, err)
+	require.Len(t, templates, 3)
+
+	p = p.WithExtensions(".yaml", ".json")
+	templates, err = p.Parse()
+	require.NoError(t, err)
+	require.Len(t, templates, 5)
+}
+
 func TestTemplateStrings(t *testing.T) {
 	tests := []struct {
 		name     string