chore: Refactor choice flags

assets/chezmoi.io/docs/reference/commands/archive.md

@@ -13,13 +13,8 @@ the extension, otherwise the default is `tar`.
 | Supported formats |
 | ----------------- |
 | `tar`             |
-| `tar.bz2`         |
 | `tar.gz`          |
-| `tar.xz`          |
-| `tar.zst`         |
-| `tbz2`            |
 | `tgz`             |
-| `txz`             |
 | `zip`             |
 
 ### `-z`, `--gzip`

internal/chezmoi/archive.go

@@ -32,52 +32,12 @@ const (
 	ArchiveFormatTarGz   ArchiveFormat = "tar.gz"
 	ArchiveFormatTarXz   ArchiveFormat = "tar.xz"
 	ArchiveFormatTarZst  ArchiveFormat = "tar.zst"
-	ArchiveFormatTbz2    ArchiveFormat = "tbz2"
-	ArchiveFormatTgz     ArchiveFormat = "tgz"
-	ArchiveFormatTxz     ArchiveFormat = "txz"
 	ArchiveFormatZip     ArchiveFormat = "zip"
 )
 
-var ArchiveFormatFlagCompletionFunc = FlagCompletionFunc([]string{
-	ArchiveFormatTar.String(),
-	ArchiveFormatTarBz2.String(),
-	ArchiveFormatTarGz.String(),
-	ArchiveFormatTarXz.String(),
-	ArchiveFormatTarZst.String(),
-	ArchiveFormatTbz2.String(),
-	ArchiveFormatTgz.String(),
-	ArchiveFormatTxz.String(),
-	ArchiveFormatZip.String(),
-})
-
-type UnknownArchiveFormatError string
-
-func (e UnknownArchiveFormatError) Error() string {
-	if e == UnknownArchiveFormatError(ArchiveFormatUnknown) {
-		return "unknown archive format"
-	}
-	return string(e) + ": unknown archive format"
-}
-
 // An WalkArchiveFunc is called once for each entry in an archive.
 type WalkArchiveFunc func(name string, info fs.FileInfo, r io.Reader, linkname string) error
 
-// Set implements github.com/spf13/pflag.Value.Set.
-func (f *ArchiveFormat) Set(s string) error {
-	*f = ArchiveFormat(s)
-	return nil
-}
-
-// String implements github.com/spf13/pflag.Value.String.
-func (f ArchiveFormat) String() string {
-	return string(f)
-}
-
-// Type implements github.com/spf13/pflag.Value.Type.
-func (f ArchiveFormat) Type() string {
-	return "format"
-}
-
 // GuessArchiveFormat guesses the archive format from the name and data.
 func GuessArchiveFormat(name string, data []byte) ArchiveFormat {
 	switch nameLower := strings.ToLower(name); {
@@ -123,17 +83,17 @@ func WalkArchive(data []byte, format ArchiveFormat, f WalkArchiveFunc) error {
 	switch format {
 	case ArchiveFormatTar:
 		// Already in tar format, do nothing.
-	case ArchiveFormatTarBz2, ArchiveFormatTbz2:
+	case ArchiveFormatTarBz2:
 		// Decompress with bzip2.
 		r = bzip2.NewReader(r)
-	case ArchiveFormatTarGz, ArchiveFormatTgz:
+	case ArchiveFormatTarGz:
 		// Decompress with gzip.
 		var err error
 		r, err = gzip.NewReader(r)
 		if err != nil {
 			return err
 		}
-	case ArchiveFormatTarXz, ArchiveFormatTxz:
+	case ArchiveFormatTarXz:
 		// Decompress with xz.
 		var err error
 		r, err = xz.NewReader(r)
@@ -148,7 +108,7 @@ func WalkArchive(data []byte, format ArchiveFormat, f WalkArchiveFunc) error {
 			return err
 		}
 	default:
-		return UnknownArchiveFormatError(format)
+		return fmt.Errorf("%s: unknown archive format", format)
 	}
 	return walkArchiveTar(r, f)
 }

internal/cmd/archivecmd.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"archive/tar"
+	"fmt"
 	"os/user"
 	"strconv"
 	"strings"
@@ -15,7 +16,7 @@ import (
 
 type archiveCmdConfig struct {
 	filter     *chezmoi.EntryTypeFilter
-	format     chezmoi.ArchiveFormat
+	format     *choiceFlag
 	gzip       bool
 	init       bool
 	parentDirs bool
@@ -37,29 +38,37 @@ func (c *Config) newArchiveCmd() *cobra.Command {
 	}
 
 	archiveCmd.Flags().VarP(c.archive.filter.Exclude, "exclude", "x", "Exclude entry types")
-	archiveCmd.Flags().VarP(&c.archive.format, "format", "f", "Set archive format")
+	archiveCmd.Flags().VarP(c.archive.format, "format", "f", "Set archive format")
+	must(archiveCmd.RegisterFlagCompletionFunc("format", c.archive.format.FlagCompletionFunc()))
 	archiveCmd.Flags().BoolVarP(&c.archive.gzip, "gzip", "z", c.archive.gzip, "Compress output with gzip")
 	archiveCmd.Flags().VarP(c.archive.filter.Exclude, "include", "i", "Include entry types")
 	archiveCmd.Flags().BoolVar(&c.archive.init, "init", c.archive.init, "Recreate config file from template")
 	archiveCmd.Flags().BoolVarP(&c.archive.parentDirs, "parent-dirs", "P", c.archive.parentDirs, "Archive parent directories")
 	archiveCmd.Flags().BoolVarP(&c.archive.recursive, "recursive", "r", c.archive.recursive, "Recurse into subdirectories")
 
-	must(archiveCmd.RegisterFlagCompletionFunc("format", chezmoi.ArchiveFormatFlagCompletionFunc))
-
 	return archiveCmd
 }
 
 func (c *Config) runArchiveCmd(cmd *cobra.Command, args []string) error {
-	format := c.archive.format
-	if format == chezmoi.ArchiveFormatUnknown {
+	var format chezmoi.ArchiveFormat
+	switch formatStr := c.archive.format.String(); formatStr {
+	case "":
 		format = chezmoi.GuessArchiveFormat(c.outputAbsPath.String(), nil)
 		if format == chezmoi.ArchiveFormatUnknown {
 			format = chezmoi.ArchiveFormatTar
 		}
+	case "tar":
+		format = chezmoi.ArchiveFormatTar
+	case "tar.gz", "tgz":
+		format = chezmoi.ArchiveFormatTarGz
+	case "zip":
+		format = chezmoi.ArchiveFormatZip
+	default:
+		return fmt.Errorf("%s: invalid format", formatStr)
 	}
 
 	gzipOutput := c.archive.gzip
-	if format == chezmoi.ArchiveFormatTarGz || format == chezmoi.ArchiveFormatTgz {
+	if format == chezmoi.ArchiveFormatTarGz {
 		gzipOutput = true
 	}
 
@@ -69,12 +78,10 @@ func (c *Config) runArchiveCmd(cmd *cobra.Command, args []string) error {
 		Close() error
 	}
 	switch format {
-	case chezmoi.ArchiveFormatTar, chezmoi.ArchiveFormatTarGz, chezmoi.ArchiveFormatTgz:
+	case chezmoi.ArchiveFormatTar, chezmoi.ArchiveFormatTarGz:
 		archiveSystem = chezmoi.NewTarWriterSystem(&output, tarHeaderTemplate())
 	case chezmoi.ArchiveFormatZip:
 		archiveSystem = chezmoi.NewZIPWriterSystem(&output, time.Now().UTC())
-	default:
-		return chezmoi.UnknownArchiveFormatError(format)
 	}
 	if err := c.applyArgs(cmd.Context(), archiveSystem, chezmoi.EmptyAbsPath, args, applyArgsOptions{
 		cmd:        cmd,

internal/cmd/choiceflag.go

@@ -0,0 +1,97 @@
+package cmd
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	"github.com/twpayne/chezmoi/v2/internal/chezmoi"
+	"github.com/twpayne/chezmoi/v2/internal/chezmoimaps"
+	"github.com/twpayne/chezmoi/v2/internal/chezmoiset"
+)
+
+// A choiceFlag is a flag which accepts a limited set of allowed values.
+type choiceFlag struct {
+	value               string
+	allowedValues       chezmoiset.Set[string]
+	uniqueAbbreviations map[string]string
+}
+
+// newChoiceFlag returns a new choiceFlag with the given value and allowed
+// values. If value is not allowed then it panics.
+func newChoiceFlag(value string, allowedValues []string) *choiceFlag {
+	allowedValuesSet := chezmoiset.New(allowedValues...)
+	if !allowedValuesSet.Contains(value) {
+		panic("value not allowed")
+	}
+	return &choiceFlag{
+		value:               value,
+		allowedValues:       allowedValuesSet,
+		uniqueAbbreviations: chezmoi.UniqueAbbreviations(allowedValues),
+	}
+}
+
+// FlagCompletionFunc returns f's flag completion function.
+func (f *choiceFlag) FlagCompletionFunc() func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
+	return chezmoi.FlagCompletionFunc(chezmoimaps.SortedKeys(f.allowedValues))
+}
+
+// MarshalJSON implements encoding/json.Marshaler.MarshalJSON.
+func (f *choiceFlag) MarshalJSON() ([]byte, error) {
+	return []byte(strconv.Quote(f.value)), nil
+}
+
+// MarshalText implements encoding.TextMarshaler.MarshalText.
+func (f *choiceFlag) MarshalText() ([]byte, error) {
+	return []byte(f.value), nil
+}
+
+// Set implements github.com/spf13/pflag.Value.Set.
+func (f *choiceFlag) Set(s string) error {
+	value, ok := f.uniqueAbbreviations[s]
+	if !ok {
+		return errors.New("invalid value")
+	}
+	f.value = value
+	return nil
+}
+
+func (f *choiceFlag) String() string {
+	return f.value
+}
+
+// Type implements github.com/spf13/pflag.Value.Type.
+func (f *choiceFlag) Type() string {
+	sortedKeys := chezmoimaps.SortedKeys(f.allowedValues)
+	if len(sortedKeys) > 0 && sortedKeys[0] == "" {
+		sortedKeys[0] = "<none>"
+	}
+	return strings.Join(sortedKeys, "|")
+}
+
+// UnmarshalJSON implements encoding/json.Unmarshaler.UnmarshalJSON.
+func (f *choiceFlag) UnmarshalJSON(data []byte) error {
+	var value string
+	if err := json.Unmarshal(data, &value); err != nil {
+		return err
+	}
+	if !f.allowedValues.Contains(value) {
+		return fmt.Errorf("%s: invalid value", value)
+	}
+	f.value = value
+	return nil
+}
+
+// UnmarshalText implements encoding.TextUnmarshaler.UnmarshalText.
+func (f *choiceFlag) UnmarshalText(text []byte) error {
+	value := string(text)
+	if !f.allowedValues.Contains(value) {
+		return fmt.Errorf("%s: invalid value", value)
+	}
+	f.value = value
+	return nil
+}