ntool/nstr/parser.go

package nstr

import (
	"errors"
	"git.noahlan.cn/noahlan/ntool/nbyte"
	"strconv"
	"strings"
	"unicode"
)

// ParseSizeOpt parse size expression options
type ParseSizeOpt struct {
	// OneAsMax if only one size value, use it as max size. default is false
	OneAsMax bool
	// SepChar is the separator char for time range string. default is '~'
	SepChar byte
	// KeywordFn is the function for parse keyword time string.
	KeywordFn func(string) (min, max uint64, err error)
}

func ensureOpt(opt *ParseSizeOpt) *ParseSizeOpt {
	if opt == nil {
		opt = &ParseSizeOpt{SepChar: '~'}
	} else {
		if opt.SepChar == 0 {
			opt.SepChar = '~'
		}
	}
	return opt
}

// ErrInvalidSizeExpr invalid size expression error
var ErrInvalidSizeExpr = errors.New("invalid size expr")

// ParseSizeRange parse range size expression to min and max size.
//
// Expression format:
//
//	"1KB~2MB"       => 1KB to 2MB
//	"-1KB"          => <1KB
//	"~1MB"          => <1MB
//	"< 1KB"         => <1KB
//	"1KB"           => >1KB
//	"1KB~"          => >1KB
//	">1KB"          => >1KB
//	"+1KB"          => >1KB
func ParseSizeRange(expr string, opt *ParseSizeOpt) (min, max uint64, err error) {
	opt = ensureOpt(opt)
	expr = strings.TrimSpace(expr)
	if expr == "" {
		err = ErrInvalidSizeExpr
		return
	}

	// parse size range. eg: "1KB~2MB"
	if strings.IndexByte(expr, '~') > -1 {
		s1, s2 := TrimCut(expr, "~")
		if s1 != "" {
			min, err = ToByteSize(s1)
			if err != nil {
				return
			}
		}

		if s2 != "" {
			max, err = ToByteSize(s2)
		}
		return
	}

	// parse single size. eg: "1KB"
	if nbyte.IsDigit(expr[0]) {
		min, err = ToByteSize(expr)
		if err != nil {
			return
		}
		if opt.OneAsMax {
			max = min
		}
		return
	}

	// parse with prefix. eg: "<1KB", ">= 1KB", "-1KB", "+1KB"
	switch expr[0] {
	case '<', '-':
		max, err = ToByteSize(strings.Trim(expr[1:], " ="))
	case '>', '+':
		min, err = ToByteSize(strings.Trim(expr[1:], " ="))
	default:
		// parse keyword. eg: "small", "large"
		if opt.KeywordFn != nil {
			min, max, err = opt.KeywordFn(expr)
		} else {
			err = ErrInvalidSizeExpr
		}
	}
	return
}

// ToByteSize converts size string like 1GB/1g or 12mb/12M into an unsigned integer number of bytes
func ToByteSize(sizeStr string) (uint64, error) {
	sizeStr = strings.TrimSpace(sizeStr)
	lastPos := len(sizeStr) - 1
	if lastPos < 0 {
		return 0, nil
	}

	if sizeStr[lastPos] == 'b' || sizeStr[lastPos] == 'B' {
		// last second char is k,m,g,t
		lastSec := sizeStr[lastPos-1]
		if lastSec > 'A' {
			lastPos--
		}
	} else if nbyte.IsDigit(sizeStr[lastPos]) { // not unit suffix. eg: 346
		return strconv.ParseUint(sizeStr, 10, 32)
	}

	multiplier := float64(1)
	switch unicode.ToLower(rune(sizeStr[lastPos])) {
	case 'k':
		multiplier = 1 << 10
	case 'm':
		multiplier = 1 << 20
	case 'g':
		multiplier = 1 << 30
	case 't':
		multiplier = 1 << 40
	case 'p':
		multiplier = 1 << 50
	default: // b
		multiplier = 1
	}

	sizeNum := strings.TrimSpace(sizeStr[:lastPos])
	size, err := strconv.ParseFloat(sizeNum, 64)
	if err != nil {
		return 0, err
	}
	return uint64(size * multiplier), nil
}