Implement Lua globals output mode

cmd/root.go

@@ -81,6 +81,7 @@ yq -P sample.json
 	rootCmd.PersistentFlags().StringVar(&yqlib.ConfiguredLuaPreferences.DocPrefix, "lua-prefix", yqlib.ConfiguredLuaPreferences.DocPrefix, "prefix")
 	rootCmd.PersistentFlags().StringVar(&yqlib.ConfiguredLuaPreferences.DocSuffix, "lua-suffix", yqlib.ConfiguredLuaPreferences.DocSuffix, "suffix")
 	rootCmd.PersistentFlags().BoolVar(&yqlib.ConfiguredLuaPreferences.UnquotedKeys, "lua-unquoted", yqlib.ConfiguredLuaPreferences.UnquotedKeys, "output unquoted string keys (e.g. {foo=\"bar\"})")
+	rootCmd.PersistentFlags().BoolVar(&yqlib.ConfiguredLuaPreferences.Globals, "lua-globals", yqlib.ConfiguredLuaPreferences.Globals, "output keys as top-level global variables")
 
 	rootCmd.PersistentFlags().BoolVarP(&nullInput, "null-input", "n", false, "Don't read input, simply evaluate the expression given. Useful for creating docs from scratch.")
 	rootCmd.PersistentFlags().BoolVarP(&noDocSeparators, "no-doc", "N", false, "Don't print document separators (---)")

pkg/yqlib/doc/usage/lua.md

@@ -57,6 +57,34 @@ return {
 };
 ```
 
+## Globals
+Uses the `--lua-globals` option to export the values into the global scope.
+
+Given a sample.yml file of:
+```yaml
+---
+country: Australia # this place
+cities:
+- Sydney
+- Melbourne
+- Brisbane
+- Perth
+```
+then
+```bash
+yq -o=lua '.' sample.yml
+```
+will output
+```lua
+country = "Australia"; -- this place
+cities = {
+	"Sydney",
+	"Melbourne",
+	"Brisbane",
+	"Perth",
+};
+```
+
 ## Elaborate example
 Given a sample.yml file of:
 ```yaml

pkg/yqlib/encoder_lua.go

@@ -14,6 +14,7 @@ type luaEncoder struct {
 	indent    int
 	indentStr string
 	unquoted  bool
+	globals   bool
 	escape    *strings.Replacer
 }
 
@@ -68,7 +69,7 @@ func NewLuaEncoder(prefs LuaPreferences) Encoder {
 		"\\t", "\t",
 		"\\\\", "\\",
 	)
-	return &luaEncoder{unescape.Replace(prefs.DocPrefix), unescape.Replace(prefs.DocSuffix), 0, "\t", prefs.UnquotedKeys, escape}
+	return &luaEncoder{unescape.Replace(prefs.DocPrefix), unescape.Replace(prefs.DocSuffix), 0, "\t", prefs.UnquotedKeys, prefs.Globals, escape}
 }
 
 func (le *luaEncoder) PrintDocumentSeparator(writer io.Writer) error {
@@ -178,16 +179,18 @@ func needsQuoting(s string) bool {
 	return false
 }
 
-func (le *luaEncoder) encodeMap(writer io.Writer, node *yaml.Node) error {
+func (le *luaEncoder) encodeMap(writer io.Writer, node *yaml.Node, global bool) error {
+	if !global {
 		err := writeString(writer, "{")
 		if err != nil {
 			return err
 		}
 		le.indent++
+	}
 	for i, child := range node.Content {
 		if (i % 2) == 1 {
 			// value
-			err = le.Encode(writer, child)
+			err := le.Encode(writer, child)
 			if err != nil {
 				return err
 			}
@@ -197,16 +200,25 @@ func (le *luaEncoder) encodeMap(writer io.Writer, node *yaml.Node) error {
 			}
 		} else {
 			// key
-			err = le.writeIndent(writer)
+			if !global || i > 0 {
+				err := le.writeIndent(writer)
 				if err != nil {
 					return err
 				}
-			if le.unquoted && child.Tag == "!!str" && !needsQuoting(child.Value) {
+			}
-				err = writeString(writer, child.Value+" = ")
+			if (le.unquoted || global) && child.Tag == "!!str" && !needsQuoting(child.Value) {
+				err := writeString(writer, child.Value+" = ")
 				if err != nil {
 					return err
 				}
 			} else {
+				if global {
+					// This only works in Lua 5.2+
+					err := writeString(writer, "_ENV")
+					if err != nil {
+						return err
+					}
+				}
 				err := writeString(writer, "[")
 				if err != nil {
 					return err
@@ -223,7 +235,7 @@ func (le *luaEncoder) encodeMap(writer io.Writer, node *yaml.Node) error {
 		}
 		if child.LineComment != "" {
 			sansPrefix, _ := strings.CutPrefix(child.LineComment, "#")
-			err = writeString(writer, strings.Repeat(" ", i%2)+"--"+sansPrefix)
+			err := writeString(writer, strings.Repeat(" ", i%2)+"--"+sansPrefix)
 			if err != nil {
 				return err
 			}
@@ -236,9 +248,12 @@ func (le *luaEncoder) encodeMap(writer io.Writer, node *yaml.Node) error {
 			}
 		}
 	}
+	if global {
+		return writeString(writer, "\n")
+	}
 	le.indent--
 	if len(node.Content) != 0 {
-		err = le.writeIndent(writer)
+		err := le.writeIndent(writer)
 		if err != nil {
 			return err
 		}
@@ -251,7 +266,7 @@ func (le *luaEncoder) encodeAny(writer io.Writer, node *yaml.Node) error {
 	case yaml.SequenceNode:
 		return le.encodeArray(writer, node)
 	case yaml.MappingNode:
-		return le.encodeMap(writer, node)
+		return le.encodeMap(writer, node, false)
 	case yaml.ScalarNode:
 		switch node.Tag {
 		case "!!str":
@@ -288,6 +303,12 @@ 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

pkg/yqlib/lua.go

@@ -4,6 +4,7 @@ type LuaPreferences struct {
 	DocPrefix    string
 	DocSuffix    string
 	UnquotedKeys bool
+	Globals      bool
 }
 
 func NewDefaultLuaPreferences() LuaPreferences {
@@ -11,6 +12,7 @@ func NewDefaultLuaPreferences() LuaPreferences {
 		DocPrefix:    "return ",
 		DocSuffix:    ";\n",
 		UnquotedKeys: false,
+		Globals:      false,
 	}
 }
 

pkg/yqlib/lua_test.go

@@ -50,6 +50,26 @@ cities:
 		"Perth",
 	};
 };
+`,
+	},
+	{
+		description:    "Globals",
+		subdescription: "Uses the `--lua-globals` option to export the values into the global scope.",
+		scenarioType:   "globals-encode",
+		input: `---
+country: Australia # this place
+cities:
+- Sydney
+- Melbourne
+- Brisbane
+- Perth`,
+		expected: `country = "Australia"; -- this place
+cities = {
+	"Sydney",
+	"Melbourne",
+	"Brisbane",
+	"Perth",
+};
 `,
 	},
 	{
@@ -155,6 +175,14 @@ func testLuaScenario(t *testing.T, s formatScenario) {
 			DocPrefix:    "return ",
 			DocSuffix:    ";\n",
 			UnquotedKeys: true,
+			Globals:      false,
+		})), s.description)
+	case "globals-encode":
+		test.AssertResultWithContext(t, s.expected, mustProcessFormatScenario(s, NewYamlDecoder(ConfiguredYamlPreferences), NewLuaEncoder(LuaPreferences{
+			DocPrefix:    "return ",
+			DocSuffix:    ";\n",
+			UnquotedKeys: false,
+			Globals:      true,
 		})), s.description)
 	default:
 		panic(fmt.Sprintf("unhandled scenario type %q", s.scenarioType))
@@ -168,7 +196,7 @@ func documentLuaScenario(t *testing.T, w *bufio.Writer, i interface{}) {
 		return
 	}
 	switch s.scenarioType {
-	case "encode", "unquoted-encode":
+	case "encode", "unquoted-encode", "globals-encode":
 		documentLuaEncodeScenario(w, s)
 	default:
 		panic(fmt.Sprintf("unhandled scenario type %q", s.scenarioType))
@@ -184,11 +212,20 @@ func documentLuaEncodeScenario(w *bufio.Writer, s formatScenario) {
 	}
 
 	prefs := ConfiguredLuaPreferences
-	if s.scenarioType == "unquoted-encode" {
+	switch s.scenarioType {
+	case "unquoted-encode":
 		prefs = LuaPreferences{
 			DocPrefix:    "return ",
 			DocSuffix:    ";\n",
 			UnquotedKeys: true,
+			Globals:      false,
+		}
+	case "globals-encode":
+		prefs = LuaPreferences{
+			DocPrefix:    "return ",
+			DocSuffix:    ";\n",
+			UnquotedKeys: false,
+			Globals:      true,
 		}
 	}
 	writeOrPanic(w, "Given a sample.yml file of:\n")
@@ -198,6 +235,8 @@ func documentLuaEncodeScenario(w *bufio.Writer, s formatScenario) {
 	switch s.scenarioType {
 	case "unquoted-encode":
 		writeOrPanic(w, "```bash\nyq -o=lua --lua-unquoted '.' sample.yml\n```\n")
+	case "globals-encode":
+		writeOrPanic(w, "```bash\nyq -o=lua --lua-globals '.' sample.yml\n```\n")
 	default:
 		writeOrPanic(w, "```bash\nyq -o=lua '.' sample.yml\n```\n")
 	}