sdk/go/arvados/byte_size.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package arvados
   6
   7 import (
   8         "encoding/json"
   9         "fmt"
  10         "math"
  11         "strings"
  12 )
  13
  14 type ByteSize int64
  15
  16 var prefixValue = map[string]int64{
  17         "":   1,
  18         "K":  1000,
  19         "Ki": 1 << 10,
  20         "M":  1000000,
  21         "Mi": 1 << 20,
  22         "G":  1000000000,
  23         "Gi": 1 << 30,
  24         "T":  1000000000000,
  25         "Ti": 1 << 40,
  26         "P":  1000000000000000,
  27         "Pi": 1 << 50,
  28         "E":  1000000000000000000,
  29         "Ei": 1 << 60,
  30 }
  31
  32 func (n *ByteSize) UnmarshalJSON(data []byte) error {
  33         if len(data) == 0 || data[0] != '"' {
  34                 var i int64
  35                 err := json.Unmarshal(data, &i)
  36                 if err != nil {
  37                         return err
  38                 }
  39                 *n = ByteSize(i)
  40                 return nil
  41         }
  42         var s string
  43         err := json.Unmarshal(data, &s)
  44         if err != nil {
  45                 return err
  46         }
  47         split := strings.LastIndexAny(s, "0123456789.+-eE") + 1
  48         if split == 0 {
  49                 return fmt.Errorf("invalid byte size %q", s)
  50         }
  51         if s[split-1] == 'E' {
  52                 // We accepted an E as if it started the exponent part
  53                 // of a json number, but if the next char isn't +, -,
  54                 // or digit, then the E must have meant Exa. Instead
  55                 // of "4.5E"+"iB" we want "4.5"+"EiB".
  56                 split--
  57         }
  58         var val json.Number
  59         dec := json.NewDecoder(strings.NewReader(s[:split]))
  60         dec.UseNumber()
  61         err = dec.Decode(&val)
  62         if err != nil {
  63                 return err
  64         }
  65         if split == len(s) {
  66                 return nil
  67         }
  68         prefix := strings.Trim(s[split:], " ")
  69         if strings.HasSuffix(prefix, "B") {
  70                 prefix = prefix[:len(prefix)-1]
  71         }
  72         pval, ok := prefixValue[prefix]
  73         if !ok {
  74                 return fmt.Errorf("invalid unit %q", strings.Trim(s[split:], " "))
  75         }
  76         if intval, err := val.Int64(); err == nil {
  77                 if pval > 1 && (intval*pval)/pval != intval {
  78                         return fmt.Errorf("size %q overflows int64", s)
  79                 }
  80                 *n = ByteSize(intval * pval)
  81                 return nil
  82         } else if floatval, err := val.Float64(); err == nil {
  83                 if floatval*float64(pval) > math.MaxInt64 {
  84                         return fmt.Errorf("size %q overflows int64", s)
  85                 }
  86                 *n = ByteSize(int64(floatval * float64(pval)))
  87                 return nil
  88         } else {
  89                 return fmt.Errorf("bug: json.Number for %q is not int64 or float64: %s", s, err)
  90         }
  91 }