github.com/keybase/client/go@v0.0.0-20240309051027-028f7c731f8b/kbfs/libkbfs/flag_size.go (about)

     1  // Copyright 2016 Keybase Inc. All rights reserved.
     2  // Use of this source code is governed by a BSD
     3  // license that can be found in the LICENSE file.
     4  
     5  package libkbfs
     6  
     7  import (
     8  	"fmt"
     9  	"strconv"
    10  	"strings"
    11  )
    12  
    13  // SizeFlag is for specifying sizes with the flag package.
    14  type SizeFlag struct {
    15  	v *int64
    16  }
    17  
    18  // Get for flag interface.
    19  func (sf SizeFlag) Get() interface{} { return *sf.v }
    20  
    21  // String for flag interface.
    22  func (sf SizeFlag) String() string {
    23  	// This happens when izZeroValue() from flag.go makes a zero
    24  	// value from the type of a flag.
    25  	if sf.v == nil {
    26  		return "0"
    27  	}
    28  	v := *sf.v
    29  	var suffix string
    30  	const units = "kmgt"
    31  	var c byte
    32  	if v%1000 == 0 {
    33  		for i := 0; v%1000 == 0 && i < len(units); i++ {
    34  			c = units[i]
    35  			v /= 1000
    36  		}
    37  		suffix = string(c)
    38  	} else if v%1024 == 0 {
    39  		for i := 0; v%1024 == 0 && i < len(units); i++ {
    40  			c = units[i]
    41  			v /= 1024
    42  		}
    43  		suffix = string(c) + "i"
    44  	}
    45  	return fmt.Sprintf("%d%s", v, suffix)
    46  }
    47  
    48  // Set for flag interface.
    49  func (sf SizeFlag) Set(raw string) error {
    50  	i := 0
    51  	for ; i < len(raw) && raw[i] >= '0' && raw[i] <= '9'; i++ {
    52  	}
    53  	val, err := strconv.ParseInt(raw[:i], 10, 64)
    54  	if err != nil {
    55  		return err
    56  	}
    57  	// constant map so concurrent access is ok.
    58  	mult, found := sizeFlagMap[strings.ToLower(raw[i:])]
    59  	if !found {
    60  		return fmt.Errorf("Invalid syntax: %q, supported syntax is 0-9+(k|m|g|t)[i]", raw)
    61  	}
    62  	*sf.v = val * mult
    63  	return nil
    64  }
    65  
    66  var sizeFlagMap = map[string]int64{
    67  	"":   1,
    68  	"k":  1000,
    69  	"ki": 1024,
    70  	"m":  1000 * 1000,
    71  	"mi": 1024 * 1024,
    72  	"g":  1000 * 1000 * 1000,
    73  	"gi": 1024 * 1024 * 1024,
    74  	"t":  1000 * 1000 * 1000 * 1000,
    75  	"ti": 1024 * 1024 * 1024 * 1024,
    76  }