1package units
2
3import (
4	"fmt"
5	"regexp"
6	"strconv"
7	"strings"
8)
9
10// See: http://en.wikipedia.org/wiki/Binary_prefix
11const (
12	// Decimal
13
14	KB = 1000
15	MB = 1000 * KB
16	GB = 1000 * MB
17	TB = 1000 * GB
18	PB = 1000 * TB
19
20	// Binary
21
22	KiB = 1024
23	MiB = 1024 * KiB
24	GiB = 1024 * MiB
25	TiB = 1024 * GiB
26	PiB = 1024 * TiB
27)
28
29type unitMap map[string]int64
30
31var (
32	decimalMap = unitMap{"k": KB, "m": MB, "g": GB, "t": TB, "p": PB}
33	binaryMap  = unitMap{"k": KiB, "m": MiB, "g": GiB, "t": TiB, "p": PiB}
34	sizeRegex  = regexp.MustCompile(`^(\d+(\.\d+)*) ?([kKmMgGtTpP])?[iI]?[bB]?$`)
35)
36
37var decimapAbbrs = []string{"B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}
38var binaryAbbrs = []string{"B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"}
39
40func getSizeAndUnit(size float64, base float64, _map []string) (float64, string) {
41	i := 0
42	unitsLimit := len(_map) - 1
43	for size >= base && i < unitsLimit {
44		size = size / base
45		i++
46	}
47	return size, _map[i]
48}
49
50// CustomSize returns a human-readable approximation of a size
51// using custom format.
52func CustomSize(format string, size float64, base float64, _map []string) string {
53	size, unit := getSizeAndUnit(size, base, _map)
54	return fmt.Sprintf(format, size, unit)
55}
56
57// HumanSizeWithPrecision allows the size to be in any precision,
58// instead of 4 digit precision used in units.HumanSize.
59func HumanSizeWithPrecision(size float64, precision int) string {
60	size, unit := getSizeAndUnit(size, 1000.0, decimapAbbrs)
61	return fmt.Sprintf("%.*g%s", precision, size, unit)
62}
63
64// HumanSize returns a human-readable approximation of a size
65// capped at 4 valid numbers (eg. "2.746 MB", "796 KB").
66func HumanSize(size float64) string {
67	return HumanSizeWithPrecision(size, 4)
68}
69
70// BytesSize returns a human-readable size in bytes, kibibytes,
71// mebibytes, gibibytes, or tebibytes (eg. "44kiB", "17MiB").
72func BytesSize(size float64) string {
73	return CustomSize("%.4g%s", size, 1024.0, binaryAbbrs)
74}
75
76// FromHumanSize returns an integer from a human-readable specification of a
77// size using SI standard (eg. "44kB", "17MB").
78func FromHumanSize(size string) (int64, error) {
79	return parseSize(size, decimalMap)
80}
81
82// RAMInBytes parses a human-readable string representing an amount of RAM
83// in bytes, kibibytes, mebibytes, gibibytes, or tebibytes and
84// returns the number of bytes, or -1 if the string is unparseable.
85// Units are case-insensitive, and the 'b' suffix is optional.
86func RAMInBytes(size string) (int64, error) {
87	return parseSize(size, binaryMap)
88}
89
90// Parses the human-readable size string into the amount it represents.
91func parseSize(sizeStr string, uMap unitMap) (int64, error) {
92	matches := sizeRegex.FindStringSubmatch(sizeStr)
93	if len(matches) != 4 {
94		return -1, fmt.Errorf("invalid size: '%s'", sizeStr)
95	}
96
97	size, err := strconv.ParseFloat(matches[1], 64)
98	if err != nil {
99		return -1, err
100	}
101
102	unitPrefix := strings.ToLower(matches[3])
103	if mul, ok := uMap[unitPrefix]; ok {
104		size *= float64(mul)
105	}
106
107	return int64(size), nil
108}
109