Merge pull request #3998 from jandubois/yq-options

Add `--yq` option to `limactl list` and `limactl info`

Author: AkihiroSuda

cmd/limactl/info.go

@@ -7,9 +7,12 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/mikefarah/yq/v4/pkg/yqlib"
 	"github.com/spf13/cobra"
 
 	"github.com/lima-vm/lima/v2/pkg/limainfo"
+	"github.com/lima-vm/lima/v2/pkg/uiutil"
+	"github.com/lima-vm/lima/v2/pkg/yqutil"
 )
 
 func newInfoCommand() *cobra.Command {
@@ -20,11 +23,19 @@ func newInfoCommand() *cobra.Command {
 		RunE:    infoAction,
 		GroupID: advancedCommand,
 	}
+	infoCommand.Flags().String("yq", ".", "Apply yq expression to output")
+
 	return infoCommand
 }
 
 func infoAction(cmd *cobra.Command, _ []string) error {
 	ctx := cmd.Context()
+
+	yq, err := cmd.Flags().GetString("yq")
+	if err != nil {
+		return err
+	}
+
 	info, err := limainfo.New(ctx)
 	if err != nil {
 		return err
@@ -33,6 +44,14 @@ func infoAction(cmd *cobra.Command, _ []string) error {
 	if err != nil {
 		return err
 	}
-	_, err = fmt.Fprintln(cmd.OutOrStdout(), string(j))
+
+	encoderPrefs := yqlib.ConfiguredJSONPreferences.Copy()
+	encoderPrefs.Indent = 4
+	encoderPrefs.ColorsEnabled = uiutil.OutputIsTTY(cmd.OutOrStdout())
+	encoder := yqlib.NewJSONEncoder(encoderPrefs)
+	str, err := yqutil.EvaluateExpressionWithEncoder(yq, string(j), encoder)
+	if err == nil {
+		_, err = fmt.Fprint(cmd.OutOrStdout(), str)
+	}
 	return err
 }

cmd/limactl/list.go

@@ -4,20 +4,23 @@
 package main
 
 import (
+	"bufio"
+	"bytes"
 	"errors"
 	"fmt"
-	"os"
 	"reflect"
 	"sort"
 	"strings"
 
 	"github.com/cheggaaa/pb/v3/termutil"
-	"github.com/mattn/go-isatty"
+	"github.com/mikefarah/yq/v4/pkg/yqlib"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 
 	"github.com/lima-vm/lima/v2/pkg/limatype"
 	"github.com/lima-vm/lima/v2/pkg/store"
+	"github.com/lima-vm/lima/v2/pkg/uiutil"
+	"github.com/lima-vm/lima/v2/pkg/yqutil"
 )
 
 func fieldNames() []string {
@@ -64,6 +67,7 @@ The following legacy flags continue to function:
 	listCommand.Flags().Bool("json", false, "JSONify output")
 	listCommand.Flags().BoolP("quiet", "q", false, "Only show names")
 	listCommand.Flags().Bool("all-fields", false, "Show all fields")
+	listCommand.Flags().String("yq", "", "Apply yq expression to each instance")
 
 	return listCommand
 }
@@ -109,6 +113,10 @@ func listAction(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+	yq, err := cmd.Flags().GetString("yq")
+	if err != nil {
+		return err
+	}
 
 	if jsonFormat {
 		format = "json"
@@ -121,6 +129,14 @@ func listAction(cmd *cobra.Command, args []string) error {
 	if listFields && cmd.Flags().Changed("format") {
 		return errors.New("option --list-fields conflicts with option --format")
 	}
+	if yq != "" {
+		if cmd.Flags().Changed("format") && format != "json" && format != "yaml" {
+			return errors.New("option --yq only works with --format json or yaml")
+		}
+		if listFields {
+			return errors.New("option --list-fields conflicts with option --yq")
+		}
+	}
 
 	if quiet && format != "table" {
 		return errors.New("option --quiet can only be used with '--format table'")
@@ -194,16 +210,80 @@ func listAction(cmd *cobra.Command, args []string) error {
 	}
 
 	options := store.PrintOptions{AllFields: allFields}
-	out := cmd.OutOrStdout()
-	if out == os.Stdout {
-		if isatty.IsTerminal(os.Stdout.Fd()) || isatty.IsCygwinTerminal(os.Stdout.Fd()) {
-			if w, err := termutil.TerminalWidth(); err == nil {
-				options.TerminalWidth = w
+	isTTY := uiutil.OutputIsTTY(cmd.OutOrStdout())
+	if isTTY {
+		if w, err := termutil.TerminalWidth(); err == nil {
+			options.TerminalWidth = w
+		}
+	}
+	// --yq implies --format json unless --format yaml has been explicitly specified
+	if yq != "" && !cmd.Flags().Changed("format") {
+		format = "json"
+	}
+	// Always pipe JSON and YAML through yq to colorize it if isTTY
+	if yq == "" && (format == "json" || format == "yaml") {
+		yq = "."
+	}
+
+	if yq == "" {
+		err = store.PrintInstances(cmd.OutOrStdout(), instances, format, &options)
+		if err == nil && unmatchedInstances {
+			return unmatchedInstancesError{}
+		}
+		return err
+	}
+
+	buf := new(bytes.Buffer)
+	err = store.PrintInstances(buf, instances, format, &options)
+	if err != nil {
+		return err
+	}
+
+	if format == "json" {
+		encoderPrefs := yqlib.ConfiguredJSONPreferences.Copy()
+		if isTTY {
+			// Using non-0 indent means the instance will be printed over multiple lines,
+			// so is no longer in JSON Lines format. This is a compromise for readability.
+			encoderPrefs.Indent = 4
+			encoderPrefs.ColorsEnabled = true
+		} else {
+			encoderPrefs.Indent = 0
+			encoderPrefs.ColorsEnabled = false
+		}
+		encoder := yqlib.NewJSONEncoder(encoderPrefs)
+
+		// Each line contains the JSON object for one Lima instance.
+		scanner := bufio.NewScanner(buf)
+		for scanner.Scan() {
+			var str string
+			if str, err = yqutil.EvaluateExpressionWithEncoder(yq, scanner.Text(), encoder); err != nil {
+				return err
 			}
+			if _, err = fmt.Fprint(cmd.OutOrStdout(), str); err != nil {
+				return err
+			}
+		}
+		err = scanner.Err()
+		if err == nil && unmatchedInstances {
+			return unmatchedInstancesError{}
 		}
+		return err
 	}
 
-	err = store.PrintInstances(out, instances, format, &options)
+	var str string
+	if isTTY {
+		// This branch is trading the better formatting from yamlfmt for colorizing from yqlib.
+		if str, err = yqutil.EvaluateExpressionPlain(yq, buf.String(), true); err != nil {
+			return err
+		}
+	} else {
+		var res []byte
+		if res, err = yqutil.EvaluateExpression(yq, buf.Bytes()); err != nil {
+			return err
+		}
+		str = string(res)
+	}
+	_, err = fmt.Fprint(cmd.OutOrStdout(), str)
 	if err == nil && unmatchedInstances {
 		return unmatchedInstancesError{}
 	}

cmd/limactl/template.go

@@ -17,6 +17,7 @@ import (
 	"github.com/lima-vm/lima/v2/pkg/limatmpl"
 	"github.com/lima-vm/lima/v2/pkg/limatype/dirnames"
 	"github.com/lima-vm/lima/v2/pkg/limayaml"
+	"github.com/lima-vm/lima/v2/pkg/uiutil"
 	"github.com/lima-vm/lima/v2/pkg/yqutil"
 )
 
@@ -148,6 +149,15 @@ func templateCopyAction(cmd *cobra.Command, args []string) error {
 			return err
 		}
 	}
+	if target == "-" && uiutil.OutputIsTTY(cmd.OutOrStdout()) {
+		// run the output through YQ to colorize it
+		out, err := yqutil.EvaluateExpressionPlain(".", string(tmpl.Bytes), true)
+		if err == nil {
+			_, err = fmt.Fprint(cmd.OutOrStdout(), out)
+		}
+		return err
+	}
+
 	writer := cmd.OutOrStdout()
 	if target != "-" {
 		file, err := os.OpenFile(target, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
@@ -200,7 +210,8 @@ func templateYQAction(cmd *cobra.Command, args []string) error {
 	if err := fillDefaults(ctx, tmpl); err != nil {
 		return err
 	}
-	out, err := yqutil.EvaluateExpressionPlain(expr, string(tmpl.Bytes))
+	colorsEnabled := uiutil.OutputIsTTY(cmd.OutOrStdout())
+	out, err := yqutil.EvaluateExpressionPlain(expr, string(tmpl.Bytes), colorsEnabled)
 	if err == nil {
 		_, err = fmt.Fprint(cmd.OutOrStdout(), out)
 	}

pkg/uiutil/uiutil.go

@@ -4,8 +4,12 @@
 package uiutil
 
 import (
+	"io"
+	"os"
+
 	"github.com/AlecAivazis/survey/v2"
 	"github.com/AlecAivazis/survey/v2/terminal"
+	"github.com/mattn/go-isatty"
 )
 
 var InterruptErr = terminal.InterruptErr
@@ -36,3 +40,14 @@ func Select(message string, options []string) (int, error) {
 	}
 	return ans, nil
 }
+
+// OutputIsTTY returns true if writer is going to stdout, and stdout is a terminal device,
+// not a regular file, stream, or pipe etc.
+func OutputIsTTY(writer io.Writer) bool {
+	// This setting is needed so we can write integration tests for the TTY output.
+	// It is probably not useful otherwise.
+	if os.Getenv("_LIMA_OUTPUT_IS_TTY") != "" {
+		return true
+	}
+	return writer == os.Stdout && (isatty.IsTerminal(os.Stdout.Fd()) || isatty.IsCygwinTerminal(os.Stdout.Fd()))
+}

pkg/yqutil/yqutil.go

@@ -39,8 +39,8 @@ func ValidateContent(content []byte) error {
 	return err
 }
 
-// EvaluateExpressionPlain evaluates the yq expression and returns the yq result.
-func EvaluateExpressionPlain(expression, content string) (string, error) {
+// EvaluateExpressionWithEncoder evaluates the yq expression and returns the yq result using a custom encoder.
+func EvaluateExpressionWithEncoder(expression, content string, encoder yqlib.Encoder) (string, error) {
 	if expression == "" {
 		return content, nil
 	}
@@ -50,10 +50,6 @@ func EvaluateExpressionPlain(expression, content string) (string, error) {
 	logging.SetBackend(backend)
 	yqlib.InitExpressionParser()
 
-	encoderPrefs := yqlib.ConfiguredYamlPreferences.Copy()
-	encoderPrefs.Indent = 2
-	encoderPrefs.ColorsEnabled = false
-	encoder := yqlib.NewYamlEncoder(encoderPrefs)
 	decoder := yqlib.NewYamlDecoder(yqlib.ConfiguredYamlPreferences)
 	out, err := yqlib.NewStringEvaluator().EvaluateAll(expression, content, encoder, decoder)
 	if err != nil {
@@ -82,6 +78,15 @@ func EvaluateExpressionPlain(expression, content string) (string, error) {
 	return out, nil
 }
 
+// EvaluateExpressionPlain evaluates the yq expression and returns the yq result.
+func EvaluateExpressionPlain(expression, content string, colorsEnabled bool) (string, error) {
+	encoderPrefs := yqlib.ConfiguredYamlPreferences.Copy()
+	encoderPrefs.Indent = 2
+	encoderPrefs.ColorsEnabled = colorsEnabled
+	encoder := yqlib.NewYamlEncoder(encoderPrefs)
+	return EvaluateExpressionWithEncoder(expression, content, encoder)
+}
+
 // EvaluateExpression evaluates the yq expression and returns the output formatted with yamlfmt.
 func EvaluateExpression(expression string, content []byte) ([]byte, error) {
 	if expression == "" {
@@ -101,7 +106,7 @@ func EvaluateExpression(expression string, content []byte) ([]byte, error) {
 		return nil, err
 	}
 
-	out, err := EvaluateExpressionPlain(expression, string(contentModified))
+	out, err := EvaluateExpressionPlain(expression, string(contentModified), false)
 	if err != nil {
 		return nil, err
 	}