unit.go

package utils

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

type StorageCapacityUnit uint64

func (u StorageCapacityUnit) String() string {
	switch u {
	case Byte:
		return "Byte"
	case KB:
		return "KB"
	case MB:
		return "MB"
	case GB:
		return "GB"
	case TB:
		return "TB"
	case PB:
		return "PB"
	case KiB:
		return "KiB"
	case MiB:
		return "MiB"
	case GiB:
		return "GiB"
	case TiB:
		return "TiB"
	case PiB:
		return "PiB"
	}
	return "unknow unit"
}

const (
	Byte StorageCapacityUnit = 1
	KB   StorageCapacityUnit = Byte * 1000
	MB   StorageCapacityUnit = KB * 1000
	GB   StorageCapacityUnit = MB * 1000
	TB   StorageCapacityUnit = GB * 1000
	PB   StorageCapacityUnit = TB * 1000

	KiB StorageCapacityUnit = 1 << 10
	MiB StorageCapacityUnit = 1 << 20
	GiB StorageCapacityUnit = 1 << 30
	TiB StorageCapacityUnit = 1 << 40
	PiB StorageCapacityUnit = 1 << 50
)

var (
	ReStorageSize = regexp.MustCompile(`^([1-9][0-9]*)((?:[KMGTPkmgtp]|)i?[bB])$`)
	ReUnit        = regexp.MustCompile(`^([kmgtpKMGTP]|)(|i)([bB])$`)
)

// MemorySize 内存大小，Num代表数字，Unit代表单位
type MemorySize struct {
	Num  uint64
	Unit StorageCapacityUnit
}

func (s MemorySize) HumanReadStr(i int) string {
	total := s.Num * uint64(s.Unit)
	units := []StorageCapacityUnit{Byte, KiB, MiB, GiB, TiB, PiB}
	var target StorageCapacityUnit
	for k, v := range units {
		if total/uint64(v) < 1 {
			target = units[k-1]
			break
		} else {
			target = v
		}
	}
	num := float64(total) / float64(target)
	switch i {
	case 0:
		return fmt.Sprintf("%d %s", int(num), target)
	case 1:
		return fmt.Sprintf("%.1f %s", num, target)
	case 2:
		return fmt.Sprintf("%.2f %s", num, target)
	case 3:
		return fmt.Sprintf("%.3f %s", num, target)
	default:
		return fmt.Sprintf("%.3f %s", num, target)
	}

}

func ParseUnit(s string) (StorageCapacityUnit, error) {
	s = strings.Trim(strings.TrimSpace(s), "\n")
	s = strings.ReplaceAll(s, " ", "")
	if !ReUnit.MatchString(s) {
		return 0, fmt.Errorf("invalid storage size unit: %s", s)
	}
	s = strings.ToLower(s)
	// [MiB M i B]
	// [KB K '' B]
	// [B '' '' B]
	matchs := ReUnit.FindStringSubmatch(s)
	if len(matchs) < 4 {
		return 0, fmt.Errorf("invalid storage size unit: %s", s)
	}
	isI := matchs[2] == "i"
	switch matchs[1] {
	case "":
		return Byte, nil
	case "k":
		if isI {
			return KiB, nil
		}
		return KB, nil
	case "m":
		if isI {
			return MiB, nil
		}
		return MB, nil
	case "g":
		if isI {
			return GiB, nil
		}
		return GB, nil
	case "t":
		if isI {
			return TiB, nil
		}
		return TB, nil
	case "p":
		if isI {
			return PiB, nil
		}
		return PB, nil
	default:
		return 0, fmt.Errorf("invalid storage size unit: %s", s)
	}
}

// ParseMemorySize 解析容量字符串，支持的格式有：123MiB 123MB "123 MiB" "123 MB"
func ParseMemorySize(s string) (*MemorySize, error) {
	s = strings.TrimSpace(strings.Trim(s, " \n"))
	s = strings.ReplaceAll(s, " ", "")
	if !ReStorageSize.MatchString(s) {
		return nil, fmt.Errorf("invalid memory size format: %s", s)
	}
	matchs := ReStorageSize.FindStringSubmatch(s)
	if len(matchs) < 3 {
		return nil, fmt.Errorf("invalid memory size format: %s", s)
	}

	num, err := strconv.ParseUint(matchs[1], 10, 64)
	if err != nil {
		return nil, err
	}
	result := MemorySize{}
	result.Num = num
	unit, err := ParseUnit(matchs[2])
	if err != nil {
		return nil, err
	}
	result.Unit = unit
	return &result, nil
}