From 8766c409e79e485b1b98cd52c319a4ad8ae5e774 Mon Sep 17 00:00:00 2001
From: Kenichi Kocha <64531971+kkocha@users.noreply.github.com>
Date: Thu, 20 Apr 2023 06:21:11 +0000
Subject: [PATCH] feat(env): add output flag to env command

Add `--output` flag to `helm env` command.
Supported formats are env(key-value), python, and yaml.

Closes #11909

Signed-off-by: Kenichi Kocha <64531971+kkocha@users.noreply.github.com>
---
 cmd/helm/env.go        | 135 ++++++++++++++++++++++++++++++++++++++---
 pkg/cli/environment.go |   2 +-
 2 files changed, 126 insertions(+), 11 deletions(-)

diff --git a/cmd/helm/env.go b/cmd/helm/env.go
index 3754b748d..487a904f5 100644
--- a/cmd/helm/env.go
+++ b/cmd/helm/env.go
@@ -19,18 +19,24 @@ package main
 import (
 	"fmt"
 	"io"
+	"log"
 	"sort"
+	"strings"
 
 	"github.com/spf13/cobra"
 
 	"helm.sh/helm/v3/cmd/helm/require"
+	"helm.sh/helm/v3/pkg/cli/output"
 )
 
 var envHelp = `
 Env prints out all the environment information in use by Helm.
 `
 
+const envOutputFlag string = "output"
+
 func newEnvCmd(out io.Writer) *cobra.Command {
+	outfmtEnv := keyValueENV
 	cmd := &cobra.Command{
 		Use:   "env",
 		Short: "helm client environment information",
@@ -44,22 +50,20 @@ func newEnvCmd(out io.Writer) *cobra.Command {
 
 			return nil, cobra.ShellCompDirectiveNoFileComp
 		},
-		Run: func(cmd *cobra.Command, args []string) {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			envVars := settings.EnvVars()
 
 			if len(args) == 0 {
-				// Sort the variables by alphabetical order.
-				// This allows for a constant output across calls to 'helm env'.
-				keys := getSortedEnvVarKeys()
-
-				for _, k := range keys {
-					fmt.Fprintf(out, "%s=\"%s\"\n", k, envVars[k])
-				}
-			} else {
-				fmt.Fprintf(out, "%s\n", envVars[args[0]])
+				return outfmtEnv.WriteEnvs(out, envVars)
 			}
+
+			key := args[0]
+			return outfmtEnv.WriteSingleEnv(out, key, envVars[key])
 		},
 	}
+
+	bindEnvOutputFlag(cmd, &outfmtEnv)
+
 	return cmd
 }
 
@@ -74,3 +78,114 @@ func getSortedEnvVarKeys() []string {
 
 	return keys
 }
+
+type envs map[string]string
+type envFormat string
+
+const (
+	keyValueENV envFormat = "env"
+	jsonENV     envFormat = "json"
+	yamlENV     envFormat = "yaml"
+)
+
+func envFormats() []string {
+	return []string{keyValueENV.String(), jsonENV.String(), yamlENV.String()}
+}
+
+func envFormatWithDesc() map[string]string {
+	return map[string]string{
+		keyValueENV.String(): "Output result in KEY=VALUE format",
+		jsonENV.String():     "Output result in JSON format",
+		yamlENV.String():     "Output result in YAML format",
+	}
+}
+
+func (o envFormat) String() string {
+	return string(o)
+}
+
+func (o *envFormat) Set(s string) error {
+	outfmt, err := parseFormat(s)
+	if err != nil {
+		return err
+	}
+	*o = outfmt
+	return nil
+}
+
+func (o envFormat) Type() string {
+	return "format"
+}
+
+func parseFormat(s string) (out envFormat, err error) {
+	switch s {
+	case keyValueENV.String():
+		out, err = keyValueENV, nil
+	case jsonENV.String():
+		out, err = jsonENV, nil
+	case yamlENV.String():
+		out, err = yamlENV, nil
+	default:
+		out, err = "", output.ErrInvalidFormatType
+	}
+	return
+}
+
+func bindEnvOutputFlag(cmd *cobra.Command, varRef *envFormat) {
+	cmd.Flags().VarP(varRef, envOutputFlag, "o",
+		fmt.Sprintf("prints the output in the specified format. Allowed values: %s", strings.Join(envFormats(), ", ")))
+
+	err := cmd.RegisterFlagCompletionFunc(envOutputFlag, func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		var formatNames []string
+		for format, desc := range envFormatWithDesc() {
+			formatNames = append(formatNames, fmt.Sprintf("%s\t%s", format, desc))
+		}
+
+		// Sort the results to get a deterministic order for the tests
+		sort.Strings(formatNames)
+		return formatNames, cobra.ShellCompDirectiveNoFileComp
+	})
+
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+func (o envFormat) WriteEnvs(out io.Writer, e envs) error {
+	switch o {
+	case keyValueENV:
+		return writeKeyValues(out, e)
+	case jsonENV:
+		return output.EncodeJSON(out, e)
+	case yamlENV:
+		return output.EncodeYAML(out, e)
+	default:
+		return output.ErrInvalidFormatType
+	}
+}
+
+func (o envFormat) WriteSingleEnv(out io.Writer, key, value string) error {
+	switch o {
+	case keyValueENV:
+		fmt.Fprintf(out, "%s\n", value)
+		return nil
+	case jsonENV:
+		return output.EncodeJSON(out, map[string]string{key: value})
+	case yamlENV:
+		return output.EncodeYAML(out, map[string]string{key: value})
+	default:
+		return output.ErrInvalidFormatType
+	}
+}
+
+func writeKeyValues(out io.Writer, e envs) error {
+	keys := getSortedEnvVarKeys()
+
+	// Sort the variables by alphabetical order.
+	// This allows for a constant output across calls to 'helm env'.
+	for _, k := range keys {
+		fmt.Fprintf(out, "%s=\"%s\"\n", k, e[k])
+	}
+	return nil
+
+}
diff --git a/pkg/cli/environment.go b/pkg/cli/environment.go
index dac2a4bc1..eb62254c2 100644
--- a/pkg/cli/environment.go
+++ b/pkg/cli/environment.go
@@ -189,7 +189,7 @@ func envCSV(name string) (ls []string) {
 
 func (s *EnvSettings) EnvVars() map[string]string {
 	envvars := map[string]string{
-		"HELM_BIN":               os.Args[0],
+		"HELM_BIN":               envOr("HELM_BIN", os.Args[0]),
 		"HELM_CACHE_HOME":        helmpath.CachePath(""),
 		"HELM_CONFIG_HOME":       helmpath.ConfigPath(""),
 		"HELM_DATA_HOME":         helmpath.DataPath(""),