Prompt for country on first run and save to config file

cmd/root.go

@@ -1,19 +1,25 @@
 package cmd
 
 import (
+	"bufio"
 	"encoding/json"
 	"fmt"
+	"io/fs"
 	"log"
 	"os"
+	"path/filepath"
+	"strings"
 
 	"github.com/Xuanwo/go-locale"
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
 
 	"github.com/thegreenwebfoundation/grid-intensity-go/ember"
 )
 
 const (
-	countryCode = "country-code"
+	configFileLocation = ".config/grid-intensity/config.yaml"
+	countryCode        = "country-code"
 )
 
 // rootCmd represents the base command when called without any subcommands
@@ -29,12 +35,7 @@ grid is greener or at locations where carbon intensity is lower.
 	grid-intensity -c BOL`,
 
 	Run: func(cmd *cobra.Command, args []string) {
-		country, err := cmd.Flags().GetString(countryCode)
-		if err != nil {
-			log.Fatal(err)
-		}
-
-		err = getGridIntensityForCountry(country)
+		err := runWithError()
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -50,30 +51,92 @@ func Execute() {
 
 func init() {
 	rootCmd.Flags().StringP(countryCode, "c", "", "Country code in ISO 3 character format")
+	viper.BindPFlag(countryCode, rootCmd.Flags().Lookup(countryCode))
 }
 
 func getGridIntensityForCountry(countryCode string) error {
-	if countryCode == "" {
-		// No country code provided so try to detect it from the user's locale.
-		tag, err := locale.Detect()
+	result, err := ember.GetGridIntensityForCountry(countryCode)
+	if err != nil {
+		return err
+	}
+
+	bytes, err := json.MarshalIndent(result, "", "\t")
+	if err != nil {
+		return err
+	}
+
+	fmt.Println(string(bytes))
+	return nil
+}
+
+// getCountryCode prompts the user to enter a country code. We try to detect
+// a country code from the user's locale but the user can enter another value.
+func getCountryCode() (string, error) {
+	tag, err := locale.Detect()
+	if err != nil {
+		return "", err
+	}
+
+	region, _ := tag.Region()
+	country := region.ISO3()
+
+	fmt.Printf("Provider %s needs an ISO country code as a region parameter.\n", ember.Provider)
+	if country != "" {
+		fmt.Printf("%s detected from your locale.\n", country)
+	}
+
+	var reader = bufio.NewReader(os.Stdin)
+	country, err = reader.ReadString('\n')
+	if err != nil {
+		return "", err
+	}
+
+	return strings.TrimSpace(country), nil
+}
+
+func runWithError() error {
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return nil
+	}
+
+	configFile := filepath.Join(homeDir, configFileLocation)
+	viper.SetConfigFile(configFile)
+
+	err = viper.ReadInConfig()
+	if _, ok := err.(*fs.PathError); ok {
+		// Create config file if it doesn't exist.
+		_, err = os.Create(configFile)
 		if err != nil {
 			return err
 		}
+	} else if err != nil {
+		return err
+	}
 
-		region, _ := tag.Region()
-		countryCode = region.ISO3()
+	country := viper.GetString(countryCode)
+	if err != nil {
+		return err
 	}
 
-	result, err := ember.GetGridIntensityForCountry(countryCode)
+	if country == "" {
+		country, err = getCountryCode()
+		if err != nil {
+			return err
+		}
+
+		viper.Set(countryCode, country)
+	}
+
+	err = getGridIntensityForCountry(country)
 	if err != nil {
 		return err
 	}
 
-	bytes, err := json.MarshalIndent(result, "", "\t")
+	err = viper.WriteConfig()
 	if err != nil {
 		return err
 	}
 
-	fmt.Println(string(bytes))
 	return nil
 }

ember/data.go

@@ -15,6 +15,8 @@ import (
 //go:embed co2-intensities-ember-2021.csv
 var emberData []byte
 
+const Provider = "ember-climate.org"
+
 type ApiClient struct {
 	data map[string]GridIntensity
 }

go.mod

@@ -5,11 +5,25 @@ go 1.17
 require (
 	github.com/Xuanwo/go-locale v1.1.0
 	github.com/spf13/cobra v1.4.0
+	github.com/spf13/viper v1.12.0
 )
 
 require (
+	github.com/fsnotify/fsnotify v1.5.4 // indirect
+	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/magiconair/properties v1.8.6 // indirect
+	github.com/mitchellh/mapstructure v1.5.0 // indirect
+	github.com/pelletier/go-toml v1.9.5 // indirect
+	github.com/pelletier/go-toml/v2 v2.0.1 // indirect
+	github.com/spf13/afero v1.8.2 // indirect
+	github.com/spf13/cast v1.5.0 // indirect
+	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	golang.org/x/sys v0.0.0-20211023085530-d6a326fbbf70 // indirect
+	github.com/subosito/gotenv v1.3.0 // indirect
+	golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a // indirect
 	golang.org/x/text v0.3.7 // indirect
+	gopkg.in/ini.v1 v1.66.4 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.0 // indirect
 )