github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/cmn/cos/size.go (about)

     1  // Package cos provides common low-level types and utilities for all aistore projects.
     2  /*
     3   * Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package cos
     6  
     7  import (
     8  	"fmt"
     9  	"strconv"
    10  	"strings"
    11  
    12  	jsoniter "github.com/json-iterator/go"
    13  )
    14  
    15  const (
    16  	UnitsIEC = "iec" // default
    17  	UnitsSI  = "si"  // NOTE: currently, SI system is CLI-only (compare with cmn/cos/size.go)
    18  	UnitsRaw = "raw"
    19  )
    20  
    21  // IEC (binary) units
    22  const (
    23  	KiB = 1024
    24  	MiB = 1024 * KiB
    25  	GiB = 1024 * MiB
    26  	TiB = 1024 * GiB
    27  )
    28  
    29  // IS (metric) units
    30  const (
    31  	KB = 1000
    32  	MB = 1000 * KB
    33  	GB = 1000 * MB
    34  	TB = 1000 * GB
    35  )
    36  
    37  var suffX = []string{"KIB", "MIB", "GIB", "TIB", "KB", "MB", "GB", "TB", "K", "M", "G", "T", "B"}
    38  
    39  /////////////
    40  // SizeIEC //
    41  /////////////
    42  
    43  // is used in cmn/config; is known*** to cmn/iter-fields parser (compare w/ duration.go)
    44  
    45  type SizeIEC int64
    46  
    47  func (siz SizeIEC) MarshalJSON() ([]byte, error) { return jsoniter.Marshal(siz.String()) }
    48  func (siz SizeIEC) String() string               { return ToSizeIEC(int64(siz), 0) }
    49  
    50  func (siz *SizeIEC) UnmarshalJSON(b []byte) (err error) {
    51  	var (
    52  		n   int64
    53  		val string
    54  	)
    55  	if err = jsoniter.Unmarshal(b, &val); err != nil {
    56  		return
    57  	}
    58  	n, err = ParseSize(val, UnitsIEC)
    59  	*siz = SizeIEC(n)
    60  	return
    61  }
    62  
    63  // (compare w/ CLI `ToSizeIS`)
    64  func ToSizeIEC(b int64, digits int) string {
    65  	switch {
    66  	case b >= TiB:
    67  		return fmt.Sprintf("%.*f%s", digits, float32(b)/float32(TiB), "TiB")
    68  	case b >= GiB:
    69  		return fmt.Sprintf("%.*f%s", digits, float32(b)/float32(GiB), "GiB")
    70  	case b >= MiB:
    71  		return fmt.Sprintf("%.*f%s", digits, float32(b)/float32(MiB), "MiB")
    72  	case b >= KiB:
    73  		return fmt.Sprintf("%.*f%s", digits, float32(b)/float32(KiB), "KiB")
    74  	default:
    75  		return fmt.Sprintf("%dB", b)
    76  	}
    77  }
    78  
    79  // when `units` arg is empty conversion is defined by the suffix
    80  func ParseSize(size, units string) (int64, error) {
    81  	if size == "" {
    82  		return 0, nil
    83  	}
    84  	// validation
    85  	if units != "" {
    86  		switch units {
    87  		case "", UnitsIEC, UnitsSI, UnitsRaw:
    88  		default:
    89  			return 0, fmt.Errorf("ParseSize %q: invalid units %q (expecting %s, %s, or %s)", size, units,
    90  				UnitsRaw, UnitsSI, UnitsIEC)
    91  		}
    92  	}
    93  	// units, more validation
    94  	var (
    95  		u      = UnitsRaw
    96  		s      = strings.ToUpper(strings.TrimSpace(size))
    97  		suffix = _suffix(s)
    98  	)
    99  	if suffix == "KIB" || suffix == "MIB" || suffix == "GIB" || suffix == "TIB" {
   100  		u = UnitsIEC
   101  		if units != "" && units != UnitsIEC {
   102  			return 0, fmt.Errorf("ParseSize %q error: %q vs %q units", size, u, units)
   103  		}
   104  	} else if suffix != "" && suffix != "B" {
   105  		u = UnitsSI
   106  		if units != "" {
   107  			if units == UnitsRaw {
   108  				return 0, fmt.Errorf("ParseSize %q error: %q vs %q units", size, u, units)
   109  			}
   110  			// NOTE: the case when units (arg) take precedence over the suffix
   111  			u = units
   112  		}
   113  	}
   114  	// trim suffix and convert
   115  	if suffix != "" {
   116  		s = strings.TrimSuffix(s, suffix)
   117  	}
   118  	switch {
   119  	case strings.IndexByte(suffix, 'K') >= 0:
   120  		return _convert(s, u, KB, KiB)
   121  	case strings.IndexByte(suffix, 'M') >= 0:
   122  		return _convert(s, u, MB, MiB)
   123  	case strings.IndexByte(suffix, 'G') >= 0:
   124  		return _convert(s, u, GB, GiB)
   125  	case strings.IndexByte(suffix, 'T') >= 0:
   126  		return _convert(s, u, TB, TiB)
   127  	default:
   128  		return _convert(s, u, 1, 1)
   129  	}
   130  }
   131  
   132  func _suffix(s string) string {
   133  	for _, suffix := range suffX {
   134  		if strings.HasSuffix(s, suffix) {
   135  			return suffix
   136  		}
   137  	}
   138  	return ""
   139  }
   140  
   141  func _convert(s, units string, mult, multIEC int64) (val int64, err error) {
   142  	if strings.IndexByte(s, '.') >= 0 {
   143  		var f float64
   144  		f, err = strconv.ParseFloat(s, 64)
   145  		if err != nil {
   146  			return
   147  		}
   148  		if units == UnitsIEC {
   149  			return int64(f * float64(multIEC)), err
   150  		}
   151  		return int64(f * float64(mult)), err
   152  	}
   153  	val, err = strconv.ParseInt(s, 10, 64)
   154  	if units == UnitsIEC {
   155  		return val * multIEC, err
   156  	}
   157  	return val * mult, err
   158  }