go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/flag/flagenum/flagenum.go (about)

     1  // Copyright 2016 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package flagenum is a utility package which facilitates implementation of
    16  // flag.Value, json.Marshaler, and json.Unmarshaler interfaces via a string-to-
    17  // value mapping.
    18  package flagenum
    19  
    20  import (
    21  	"encoding/json"
    22  	"fmt"
    23  	"reflect"
    24  	"sort"
    25  	"strings"
    26  )
    27  
    28  // Enum is a mapping of enumeration key strings to values that can be used as
    29  // flags.
    30  //
    31  // Strings can be mapped to any value type that is comparable via
    32  // reflect.DeepEqual.
    33  type Enum map[string]any
    34  
    35  // GetKey performs reverse lookup of the enumeration value, returning the
    36  // key that corresponds to the value.
    37  //
    38  // If multiple keys correspond to the same value, the result is undefined.
    39  func (e Enum) GetKey(value any) string {
    40  	for k, v := range e {
    41  		if reflect.DeepEqual(v, value) {
    42  			return k
    43  		}
    44  	}
    45  	return ""
    46  }
    47  
    48  // GetValue returns the mapped enumeration value associated with a key.
    49  func (e Enum) GetValue(key string) (any, error) {
    50  	for k, v := range e {
    51  		if k == key {
    52  			return v, nil
    53  		}
    54  	}
    55  	return nil, fmt.Errorf("flagenum: Invalid value; must be one of [%s]", e.Choices())
    56  }
    57  
    58  // Choices returns a comma-separated string listing sorted enumeration choices.
    59  func (e Enum) Choices() string {
    60  	keys := make([]string, 0, len(e))
    61  	for k := range e {
    62  		if k == "" {
    63  			k = `""`
    64  		}
    65  		keys = append(keys, k)
    66  	}
    67  	sort.Strings(keys)
    68  	return strings.Join(keys, ", ")
    69  }
    70  
    71  // Sets the value v to the enumeration value mapped to the supplied key.
    72  func (e Enum) setValue(v any, key string) error {
    73  	i, err := e.GetValue(key)
    74  	if err != nil {
    75  		return err
    76  	}
    77  
    78  	vValue := reflect.ValueOf(v).Elem()
    79  	if !vValue.CanSet() {
    80  		panic(fmt.Errorf("flagenum: Cannot set supplied value, %v", vValue))
    81  	}
    82  
    83  	iValue := reflect.ValueOf(i)
    84  	if !vValue.Type().AssignableTo(iValue.Type()) {
    85  		panic(fmt.Errorf("flagenum: Enumeration type (%v) is incompatible with supplied value (%v)",
    86  			vValue.Type(), iValue.Type()))
    87  	}
    88  
    89  	vValue.Set(iValue)
    90  	return nil
    91  }
    92  
    93  // FlagSet implements flag.Value's Set semantics. It identifies the mapped value
    94  // associated with the supplied key and stores it in the supplied interface.
    95  //
    96  // The interface, v, must be a valid pointer to the mapped enumeration type.
    97  func (e Enum) FlagSet(v any, key string) error {
    98  	return e.setValue(v, key)
    99  }
   100  
   101  // FlagString implements flag.Value's String semantics.
   102  func (e Enum) FlagString(v any) string {
   103  	return e.GetKey(v)
   104  }
   105  
   106  // JSONUnmarshal implements json.Unmarshaler's UnmarshalJSON semantics. It
   107  // parses data containing a quoted string, identifies the enumeration value
   108  // associated with that string, and stores it in the supplied interface.
   109  //
   110  // The interface, v, must be a valid pointer to the mapped enumeration type.
   111  // a string corresponding to one of the enum's keys.
   112  func (e Enum) JSONUnmarshal(v any, data []byte) error {
   113  	s := ""
   114  	if err := json.Unmarshal(data, &s); err != nil {
   115  		return err
   116  	}
   117  	return e.FlagSet(v, s)
   118  }
   119  
   120  // JSONMarshal implements json.Marshaler's MarshalJSON semantics. It marshals
   121  // the value in the supplied interface to its associated key and emits a quoted
   122  // string containing that key.
   123  //
   124  // The interface, v, must be a valid pointer to the mapped enumeration type.
   125  func (e Enum) JSONMarshal(v any) ([]byte, error) {
   126  	key := e.GetKey(v)
   127  	data, err := json.Marshal(&key)
   128  	if err != nil {
   129  		return nil, err
   130  	}
   131  	return data, nil
   132  }