From 7c78a15b235b817f3895e707b95928417fd832da Mon Sep 17 00:00:00 2001
From: Kim Alvefur <zash@zash.se>
Date: Thu, 12 Oct 2023 02:25:26 +0200
Subject: [PATCH] Lua output fixes (#1811)

* encoder_lua: Handle explicitly positive infinity

* 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.

* encoder_lua: Improve Tag handling robustness

Using the method call seems more reliable in case the input parser
forgets to set the tag.
---
 acceptance_tests/output-format.sh | 41 +++++++++++++++++++++++++-
 pkg/yqlib/doc/usage/lua.md        |  4 +++
 pkg/yqlib/encoder_lua.go          | 49 ++++++++++++++++++-------------
 pkg/yqlib/lua_test.go             |  4 +++
 4 files changed, 76 insertions(+), 22 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/doc/usage/lua.md b/pkg/yqlib/doc/usage/lua.md
index d6f067e4..6b6dcc8b 100644
--- a/pkg/yqlib/doc/usage/lua.md
+++ b/pkg/yqlib/doc/usage/lua.md
@@ -129,6 +129,8 @@ numbers:
   - octal: 0o30
   - float: 123.45
   - infinity: .inf
+    plus_infinity: +.inf
+    minus_infinity: -.inf
   - not: .nan
 
 ```
@@ -162,6 +164,8 @@ return {
 		},
 		{
 			["infinity"] = (1/0);
+			["plus_infinity"] = (1/0);
+			["minus_infinity"] = (-1/0);
 		},
 		{
 			["not"] = (0/0);
diff --git a/pkg/yqlib/encoder_lua.go b/pkg/yqlib/encoder_lua.go
index d1bb81aa..c1ac1574 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
 			}
@@ -270,7 +270,7 @@ func (le *luaEncoder) encodeAny(writer io.Writer, node *yaml.Node) error {
 	case yaml.MappingNode:
 		return le.encodeMap(writer, node, false)
 	case yaml.ScalarNode:
-		switch node.Tag {
+		switch node.ShortTag() {
 		case "!!str":
 			return le.encodeString(writer, node)
 		case "!!null":
@@ -292,7 +292,7 @@ func (le *luaEncoder) encodeAny(writer io.Writer, node *yaml.Node) error {
 			return writeString(writer, strings.ToLower(node.Value))
 		case "!!float":
 			switch strings.ToLower(node.Value) {
-			case ".inf":
+			case ".inf", "+.inf":
 				return writeString(writer, "(1/0)")
 			case "-.inf":
 				return writeString(writer, "(-1/0)")
@@ -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)
 }
diff --git a/pkg/yqlib/lua_test.go b/pkg/yqlib/lua_test.go
index ddcbd023..bb36e9cc 100644
--- a/pkg/yqlib/lua_test.go
+++ b/pkg/yqlib/lua_test.go
@@ -135,6 +135,8 @@ numbers:
   - octal: 0o30
   - float: 123.45
   - infinity: .inf
+    plus_infinity: +.inf
+    minus_infinity: -.inf
   - not: .nan
 `,
 		expected: `return {
@@ -161,6 +163,8 @@ numbers:
 		},
 		{
 			["infinity"] = (1/0);
+			["plus_infinity"] = (1/0);
+			["minus_infinity"] = (-1/0);
 		},
 		{
 			["not"] = (0/0);