From e709a583f42079ea65552f713ecd66450bb7c92d Mon Sep 17 00:00:00 2001
From: Kim Alvefur <zash@zash.se>
Date: Wed, 4 Oct 2023 12:28:47 +0200
Subject: [PATCH] encoder_lua: Fix inclusion of pre-/suffix when prettyPrinted

It seems certain operations like --prettyPrint or subset selections does
not produce a DocumentNode, which is where the lua pre- and suffix was
printed, causing those to be omitted.
---
 acceptance_tests/output-format.sh | 41 +++++++++++++++++++++++++++-
 pkg/yqlib/encoder_lua.go          | 45 ++++++++++++++++++-------------
 2 files changed, 66 insertions(+), 20 deletions(-)

diff --git a/acceptance_tests/output-format.sh b/acceptance_tests/output-format.sh
index 268185d4..4e9ad325 100755
--- a/acceptance_tests/output-format.sh
+++ b/acceptance_tests/output-format.sh
@@ -271,4 +271,43 @@ EOM
   assertEquals "$expected" "$X"
 }
 
-source ./scripts/shunit2
\ No newline at end of file
+testLuaOutputPretty() {
+  cat >test.yml <<EOL
+animals:
+  cat: meow
+EOL
+
+  read -r -d '' expected << EOM
+return {
+	["animals"] = {
+		["cat"] = "meow";
+	};
+};
+EOM
+
+  X=$(./yq e --output-format=lua test.yml)
+  assertEquals "$expected" "$X"
+
+  X=$(./yq e --output-format=lua --prettyPrint test.yml)
+  assertEquals "$expected" "$X"
+
+}
+
+testLuaOutputSubset() {
+  cat >test.yml <<EOL
+animals:
+  cat: meow
+EOL
+
+  read -r -d '' expected << EOM
+return {
+	["cat"] = "meow";
+};
+EOM
+
+  X=$(./yq e --output-format=lua '.animals' test.yml)
+  assertEquals "$expected" "$X"
+
+}
+
+source ./scripts/shunit2
diff --git a/pkg/yqlib/encoder_lua.go b/pkg/yqlib/encoder_lua.go
index 8efa1889..e83687d8 100644
--- a/pkg/yqlib/encoder_lua.go
+++ b/pkg/yqlib/encoder_lua.go
@@ -129,7 +129,7 @@ func (le *luaEncoder) encodeArray(writer io.Writer, node *yaml.Node) error {
 		if err != nil {
 			return err
 		}
-		err := le.Encode(writer, child)
+		err := le.encodeAny(writer, child)
 		if err != nil {
 			return err
 		}
@@ -192,7 +192,7 @@ func (le *luaEncoder) encodeMap(writer io.Writer, node *yaml.Node, global bool)
 	for i, child := range node.Content {
 		if (i % 2) == 1 {
 			// value
-			err := le.Encode(writer, child)
+			err := le.encodeAny(writer, child)
 			if err != nil {
 				return err
 			}
@@ -305,26 +305,33 @@ func (le *luaEncoder) encodeAny(writer io.Writer, node *yaml.Node) error {
 			return fmt.Errorf("Lua encoder NYI -- %s", node.ShortTag())
 		}
 	case yaml.DocumentNode:
-		if le.globals {
-			if node.Content[0].Kind != yaml.MappingNode {
-				return fmt.Errorf("--lua-global requires a top level MappingNode")
-			}
-			return le.encodeMap(writer, node.Content[0], true)
-		}
-		err := writeString(writer, le.docPrefix)
-		if err != nil {
-			return err
-		}
-		err = le.encodeAny(writer, node.Content[0])
-		if err != nil {
-			return err
-		}
-		return writeString(writer, le.docSuffix)
+		return le.encodeAny(writer, node.Content[0])
 	default:
 		return fmt.Errorf("Lua encoder NYI -- %s", node.ShortTag())
 	}
 }
 
-func (le *luaEncoder) Encode(writer io.Writer, node *yaml.Node) error {
-	return le.encodeAny(writer, node)
+func (le *luaEncoder) encodeTopLevel(writer io.Writer, node *yaml.Node) error {
+	err := writeString(writer, le.docPrefix)
+	if err != nil {
+		return err
+	}
+	err = le.encodeAny(writer, node)
+	if err != nil {
+		return err
+	}
+	return writeString(writer, le.docSuffix)
+}
+
+func (le *luaEncoder) Encode(writer io.Writer, node *yaml.Node) error {
+	if node.Kind == yaml.DocumentNode {
+		return le.Encode(writer, node.Content[0])
+	}
+	if le.globals {
+		if node.Kind != yaml.MappingNode {
+			return fmt.Errorf("--lua-global requires a top level MappingNode")
+		}
+		return le.encodeMap(writer, node, true)
+	}
+	return le.encodeTopLevel(writer, node)
 }