unit.go

package utils

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

type StorageCapacityUnit uint64

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 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 matchs == nil || 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 matchs == nil || 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
}