codeberg.org/gruf/go-bytesize@v1.0.2/bytesize.go (about)

     1  package bytesize
     2  
     3  import (
     4  	"errors"
     5  	"math/bits"
     6  	_ "strconv"
     7  	"unsafe"
     8  )
     9  
    10  const (
    11  	// SI units
    12  	KB Size = 1e3
    13  	MB Size = 1e6
    14  	GB Size = 1e9
    15  	TB Size = 1e12
    16  	PB Size = 1e15
    17  	EB Size = 1e18
    18  
    19  	// IEC units
    20  	KiB Size = 1024
    21  	MiB Size = KiB * 1024
    22  	GiB Size = MiB * 1024
    23  	TiB Size = GiB * 1024
    24  	PiB Size = TiB * 1024
    25  	EiB Size = PiB * 1024
    26  )
    27  
    28  var (
    29  	// ErrInvalidUnit is returned when an invalid IEC/SI is provided.
    30  	ErrInvalidUnit = errors.New("bytesize: invalid unit")
    31  
    32  	// ErrInvalidFormat is returned when an invalid size value is provided.
    33  	ErrInvalidFormat = errors.New("bytesize: invalid format")
    34  
    35  	// iecpows is a precomputed table of 1024^n.
    36  	iecpows = [...]float64{
    37  		float64(KiB),
    38  		float64(MiB),
    39  		float64(GiB),
    40  		float64(TiB),
    41  		float64(PiB),
    42  		float64(EiB),
    43  	}
    44  
    45  	// sipows is a precomputed table of 1000^n.
    46  	sipows = [...]float64{
    47  		float64(KB),
    48  		float64(MB),
    49  		float64(GB),
    50  		float64(TB),
    51  		float64(PB),
    52  		float64(EB),
    53  	}
    54  
    55  	// bvals is a precomputed table of IEC unit values.
    56  	iecvals = [...]float64{
    57  		'k': float64(KiB),
    58  		'K': float64(KiB),
    59  		'M': float64(MiB),
    60  		'G': float64(GiB),
    61  		'T': float64(TiB),
    62  		'P': float64(PiB),
    63  		'E': float64(EiB),
    64  	}
    65  
    66  	// sivals is a precomputed table of SI unit values.
    67  	sivals = [...]float64{
    68  		// ASCII numbers _aren't_ valid SI unit values,
    69  		// BUT if the space containing a possible unit
    70  		// char is checked with this table -- it is valid
    71  		// to provide no unit char so unit=1 works.
    72  		'0': 1,
    73  		'1': 1,
    74  		'2': 1,
    75  		'3': 1,
    76  		'4': 1,
    77  		'5': 1,
    78  		'6': 1,
    79  		'7': 1,
    80  		'8': 1,
    81  		'9': 1,
    82  
    83  		'k': float64(KB),
    84  		'K': float64(KB),
    85  		'M': float64(MB),
    86  		'G': float64(GB),
    87  		'T': float64(TB),
    88  		'P': float64(PB),
    89  		'E': float64(EB),
    90  	}
    91  )
    92  
    93  // Size is a casting for uint64 types that provides formatting
    94  // methods for byte sizes in both IEC and SI units.
    95  type Size uint64
    96  
    97  // ParseSize will parse a valid Size from given string. Both IEC and SI units are supported.
    98  func ParseSize(s string) (Size, error) {
    99  	// Parse units from string
   100  	unit, l, err := parseUnit(s)
   101  	if err != nil {
   102  		return 0, err
   103  	}
   104  
   105  	// Parse remaining string as float
   106  	f, n, err := atof64(s[:l])
   107  	if err != nil || n != l {
   108  		return 0, ErrInvalidFormat
   109  	}
   110  
   111  	return Size(f * unit), nil
   112  }
   113  
   114  // Set implements flag.Value{}.
   115  func (sz *Size) Set(in string) error {
   116  	s, err := ParseSize(in)
   117  	if err != nil {
   118  		return err
   119  	}
   120  	*sz = s
   121  	return nil
   122  }
   123  
   124  // MarshalText implements encoding.TextMarshaler{}.
   125  func (sz *Size) MarshalText() ([]byte, error) {
   126  	const maxLen = 7 // max IEC string length
   127  	return sz.AppendFormatIEC(make([]byte, 0, maxLen)), nil
   128  }
   129  
   130  // UnmarshalText implements encoding.TextUnmarshaler{}.
   131  func (sz *Size) UnmarshalText(text []byte) error {
   132  	return sz.Set(*(*string)(unsafe.Pointer(&text)))
   133  }
   134  
   135  // AppendFormat defaults to using Size.AppendFormatIEC().
   136  func (sz Size) AppendFormat(dst []byte) []byte {
   137  	return sz.AppendFormatIEC(dst) // default
   138  }
   139  
   140  // AppendFormatSI will append SI formatted size to 'dst'.
   141  func (sz Size) AppendFormatSI(dst []byte) []byte {
   142  	if uint64(sz) < 1000 {
   143  		dst = itoa(dst, uint64(sz))
   144  		dst = append(dst, 'B')
   145  		return dst
   146  	} // above is fast-path, .appendFormat() is outlined
   147  	return sz.appendFormat(dst, 1000, &sipows, "B")
   148  }
   149  
   150  // AppendFormatIEC will append IEC formatted size to 'dst'.
   151  func (sz Size) AppendFormatIEC(dst []byte) []byte {
   152  	if uint64(sz) < 1024 {
   153  		dst = itoa(dst, uint64(sz))
   154  		dst = append(dst, 'B')
   155  		return dst
   156  	} // above is fast-path, .appendFormat() is outlined
   157  	return sz.appendFormat(dst, 1024, &iecpows, "iB")
   158  }
   159  
   160  // appendFormat will append formatted Size to 'dst', depending on base, powers table and single unit suffix.
   161  func (sz Size) appendFormat(dst []byte, base uint64, pows *[6]float64, sunit string) []byte {
   162  	const (
   163  		// min "small" unit threshold
   164  		min = 0.75
   165  
   166  		// binary unit chars.
   167  		units = `kMGTPE`
   168  	)
   169  
   170  	// Larger number: get value of
   171  	// i / unit size. We have a 'min'
   172  	// threshold after which we prefer
   173  	// using the unit 1 down
   174  	n := bits.Len64(uint64(sz)) / 10
   175  	f := float64(sz) / pows[n-1]
   176  	if f < min {
   177  		f *= float64(base)
   178  		n--
   179  	}
   180  
   181  	// Append formatted float with units
   182  	dst = ftoa(dst, f)
   183  	dst = append(dst, units[n-1])
   184  	dst = append(dst, sunit...)
   185  	return dst
   186  }
   187  
   188  // StringSI returns an SI unit string format of Size.
   189  func (sz Size) StringSI() string {
   190  	const maxLen = 6 // max SI string length
   191  	b := sz.AppendFormatSI(make([]byte, 0, maxLen))
   192  	return *(*string)(unsafe.Pointer(&b))
   193  }
   194  
   195  // StringIEC returns an IEC unit string format of Size.
   196  func (sz Size) StringIEC() string {
   197  	const maxLen = 7 // max IEC string length
   198  	b := sz.AppendFormatIEC(make([]byte, 0, maxLen))
   199  	return *(*string)(unsafe.Pointer(&b))
   200  }
   201  
   202  // String returns a string format of Size, defaults to IEC unit format.
   203  func (sz Size) String() string {
   204  	return sz.StringIEC()
   205  }
   206  
   207  // parseUnit will parse the byte size unit from string 's'.
   208  func parseUnit(s string) (float64, int, error) {
   209  	// Check for string
   210  	if len(s) < 1 {
   211  		return 0, 0, ErrInvalidFormat
   212  	}
   213  
   214  	// Strip 'byte' unit suffix
   215  	if l := len(s) - 1; s[l] == 'B' {
   216  		s = s[:l]
   217  
   218  		if len(s) < 1 {
   219  			// No remaining str before unit suffix
   220  			return 0, 0, ErrInvalidFormat
   221  		}
   222  	}
   223  
   224  	// Strip IEC binary unit suffix
   225  	if l := len(s) - 1; s[l] == 'i' {
   226  		s = s[:l]
   227  
   228  		if len(s) < 1 {
   229  			// No remaining str before unit suffix
   230  			return 0, 0, ErrInvalidFormat
   231  		}
   232  
   233  		// Location of unit char.
   234  		l := len(s) - 1
   235  		c := int(s[l])
   236  
   237  		// Check valid unit char was provided
   238  		if len(iecvals) < c || iecvals[c] == 0 {
   239  			return 0, 0, ErrInvalidUnit
   240  		}
   241  
   242  		// Return parsed IEC unit size
   243  		return iecvals[c], l, nil
   244  	}
   245  
   246  	// Location of unit char.
   247  	l := len(s) - 1
   248  	c := int(s[l])
   249  
   250  	switch {
   251  	// Check valid unit char provided
   252  	case len(sivals) < c || sivals[c] == 0:
   253  		return 0, 0, ErrInvalidUnit
   254  
   255  	// No unit char (only ascii number)
   256  	case sivals[c] == 1:
   257  		l++
   258  	}
   259  
   260  	// Return parsed SI unit size
   261  	return sivals[c], l, nil
   262  }
   263  
   264  // ftoa appends string formatted 'f' to 'dst', assumed < ~800.
   265  func ftoa(dst []byte, f float64) []byte {
   266  	switch i := uint64(f); {
   267  	// Append with 2 d.p.
   268  	case i < 10:
   269  		f *= 10
   270  
   271  		// Calculate next dec. value
   272  		d1 := uint8(uint64(f) % 10)
   273  
   274  		f *= 10
   275  
   276  		// Calculate next dec. value
   277  		d2 := uint8(uint64(f) % 10)
   278  
   279  		// Round the final value
   280  		if uint64(f*10)%10 > 4 {
   281  			d2++
   282  
   283  			// Overflow, incr 'd1'
   284  			if d2 == 10 {
   285  				d2 = 0
   286  				d1++
   287  
   288  				// Overflow, incr 'i'
   289  				if d1 == 10 {
   290  					d1 = 0
   291  					i++
   292  				}
   293  			}
   294  		}
   295  
   296  		// Append decimal value
   297  		dst = itoa(dst, i)
   298  		dst = append(dst,
   299  			'.',
   300  			'0'+d1,
   301  			'0'+d2,
   302  		)
   303  
   304  	// Append with 1 d.p.
   305  	case i < 100:
   306  		f *= 10
   307  
   308  		// Calculate next dec. value
   309  		d1 := uint8(uint64(f) % 10)
   310  
   311  		// Round the final value
   312  		if uint64(f*10)%10 > 4 {
   313  			d1++
   314  
   315  			// Overflow, incr 'i'
   316  			if d1 == 10 {
   317  				d1 = 0
   318  				i++
   319  			}
   320  		}
   321  
   322  		// Append decimal value
   323  		dst = itoa(dst, i)
   324  		dst = append(dst, '.', '0'+d1)
   325  
   326  	// No decimal places
   327  	default:
   328  		dst = itoa(dst, i)
   329  	}
   330  
   331  	return dst
   332  }
   333  
   334  // itoa appends string formatted 'i' to 'dst'.
   335  func itoa(dst []byte, i uint64) []byte {
   336  	// Assemble uint in reverse order.
   337  	var b [4]byte
   338  	bp := len(b) - 1
   339  
   340  	// Append integer
   341  	for i >= 10 {
   342  		q := i / 10
   343  		b[bp] = byte('0' + i - q*10)
   344  		bp--
   345  		i = q
   346  	} // i < 10
   347  	b[bp] = byte('0' + i)
   348  
   349  	return append(dst, b[bp:]...)
   350  }
   351  
   352  // We use the following internal strconv function usually
   353  // used internally to parse float values, as we know that
   354  // are value passed will always be of 64bit type, and knowing
   355  // the returned float string length is very helpful!
   356  //
   357  //go:linkname atof64 strconv.atof64
   358  func atof64(string) (float64, int, error)