~egtann/binp

binp/datasize.go -rw-r--r-- 1.1 KiB
7746a19a — Evan Tann fix net name in test services file 5 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
package binp

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
)

// datasize is a representation of data in kilobytes.
type datasize int

// UnmarshalJSON from a string form such as "10 KB" or "10KB". Valid datatypes
// are MB, GB, and TB.
func (d *datasize) UnmarshalJSON(byt []byte) error {
	var s string
	if err := json.Unmarshal(byt, &s); err != nil {
		return fmt.Errorf("unmarshal: %w", err)
	}

	var mbMultiplier int
	switch {
	case strings.HasSuffix(s, "MB"):
		s = strings.TrimSuffix(s, "MB")
		mbMultiplier = 1
	case strings.HasSuffix(s, "GB"):
		s = strings.TrimSuffix(s, "GB")
		mbMultiplier = 1024
	case strings.HasSuffix(s, "TB"):
		s = strings.TrimSuffix(s, "TB")
		mbMultiplier = 1024 * 1024
	}
	s = strings.TrimSpace(s)

	size, err := strconv.Atoi(s)
	if err != nil {
		return fmt.Errorf("atoi: %w", err)
	}
	*d = datasize(size * mbMultiplier)
	return nil
}

func (d datasize) String() string {
	if d < 1024 {
		// MB
		return fmt.Sprintf("%d MB", d)
	}
	if d < 1024*1024 {
		// GB
		return fmt.Sprintf("%d GB", d/1024)
	}
	if d < 1024*1024*1024 {
		// TB
		return fmt.Sprintf("%d TB", d/1024/1024)
	}
	return "invalid datasize"
}