github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/flagext/bytesize.go (about)

     1  package flagext
     2  
     3  import (
     4  	"encoding/json"
     5  	"strings"
     6  
     7  	"github.com/c2h5oh/datasize"
     8  )
     9  
    10  // ByteSize is a flag parsing compatibility type for constructing human friendly sizes.
    11  // It implements flag.Value & flag.Getter.
    12  type ByteSize uint64
    13  
    14  func (bs ByteSize) String() string {
    15  	return datasize.ByteSize(bs).String()
    16  }
    17  
    18  func (bs *ByteSize) Set(s string) error {
    19  	var v datasize.ByteSize
    20  
    21  	// Bytesize currently doesn't handle things like Mb, but only handles MB.
    22  	// Therefore we capitalize just for convenience
    23  	if err := v.UnmarshalText([]byte(strings.ToUpper(s))); err != nil {
    24  		return err
    25  	}
    26  	*bs = ByteSize(v.Bytes())
    27  	return nil
    28  }
    29  
    30  func (bs ByteSize) Get() interface{} {
    31  	return bs.Val()
    32  }
    33  
    34  func (bs ByteSize) Val() int {
    35  	return int(bs)
    36  }
    37  
    38  /// UnmarshalYAML the Unmarshaler interface of the yaml pkg.
    39  func (bs *ByteSize) UnmarshalYAML(unmarshal func(interface{}) error) error {
    40  	var str string
    41  	err := unmarshal(&str)
    42  	if err != nil {
    43  		return err
    44  	}
    45  
    46  	return bs.Set(str)
    47  }
    48  
    49  // MarshalYAML implements yaml.Marshaller.
    50  // Use a string representation for consistency
    51  func (bs *ByteSize) MarshalYAML() (interface{}, error) {
    52  	return bs.String(), nil
    53  }
    54  
    55  // UnmarshalJSON implements json.Unmarsal interface to work with JSON.
    56  func (bs *ByteSize) UnmarshalJSON(val []byte) error {
    57  	var str string
    58  
    59  	if err := json.Unmarshal(val, &str); err != nil {
    60  		return err
    61  	}
    62  
    63  	return bs.Set(str)
    64  }
    65  
    66  // Use a string representation for consistency
    67  func (bs *ByteSize) MarshalJSON() ([]byte, error) {
    68  	return json.Marshal(bs.String())
    69  }