github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/util/bytesize/bytesize.go (about)

     1  package bytesize
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"regexp"
     7  	"strconv"
     8  	"strings"
     9  )
    10  
    11  type ByteSize int64
    12  
    13  const Byte ByteSize = 1
    14  
    15  const (
    16  	// TODO: fix units - SI and IEC standards swapped.
    17  
    18  	KB = 1024 * Byte
    19  	MB = 1024 * KB
    20  	GB = 1024 * MB
    21  	TB = 1024 * GB
    22  	PB = 1024 * TB
    23  
    24  	KiB = 1000 * Byte
    25  	MiB = 1000 * KiB
    26  	GiB = 1000 * MiB
    27  	TiB = 1000 * GiB
    28  	PiB = 1000 * TiB
    29  )
    30  
    31  var suffixes = []string{"KB", "MB", "GB", "TB", "PB"}
    32  
    33  func (b ByteSize) String() string {
    34  	if b < KB {
    35  		return fmt.Sprintf("%d bytes", b)
    36  	}
    37  	bf := float64(b)
    38  	for _, s := range suffixes {
    39  		bf /= 1024.0
    40  		if bf < 1024 {
    41  			return fmt.Sprintf("%.2f %s", bf, s)
    42  		}
    43  	}
    44  	return fmt.Sprintf("%.2f %s", bf, suffixes[len(suffixes)-1])
    45  }
    46  
    47  func (b ByteSize) Bytes() int { return int(b) }
    48  
    49  var multipliers = map[string]ByteSize{
    50  	"":    Byte,
    51  	"b":   Byte,
    52  	"kb":  KB,
    53  	"mb":  MB,
    54  	"gb":  GB,
    55  	"tb":  TB,
    56  	"pb":  PB,
    57  	"kib": KiB,
    58  	"mib": MiB,
    59  	"gib": GiB,
    60  	"tib": TiB,
    61  	"pib": PiB,
    62  }
    63  
    64  var byteSizeRegexp = regexp.MustCompile("^([\\d\\.]+)\\s*([^\\d]*)$")
    65  
    66  var errParse = errors.New("could not parse ByteSize")
    67  
    68  func Parse(str string) (ByteSize, error) {
    69  	r := byteSizeRegexp.FindStringSubmatch(strings.TrimSpace(str))
    70  	if len(r) != 3 {
    71  		return 0, errParse
    72  	}
    73  
    74  	multiplier, ok := multipliers[strings.ToLower(r[2])]
    75  	if !ok {
    76  		return 0, errParse
    77  	}
    78  
    79  	if strings.Contains(r[1], ".") {
    80  		val, err := strconv.ParseFloat(r[1], 64)
    81  		if err != nil {
    82  			return 0, errParse
    83  		}
    84  		return ByteSize(val * float64(multiplier)), nil
    85  	}
    86  
    87  	val, err := strconv.ParseUint(r[1], 10, 64)
    88  	if err != nil {
    89  		return 0, errParse
    90  	}
    91  	return ByteSize(val) * multiplier, nil
    92  }
    93  
    94  func (b *ByteSize) Set(value string) error {
    95  	v, err := Parse(value)
    96  	if err != nil {
    97  		return err
    98  	}
    99  	*b = v
   100  	return nil
   101  }