package config import ( "fmt" "github.com/dustin/go-humanize" ) // Size represents a byte size that can be specified in configuration files. // It can unmarshal from both numeric values (interpreted as bytes) and // human-readable strings like "10MB", "2.5GB", or "1TB". type Size int64 // UnmarshalYAML implements yaml.Unmarshaler for Size, allowing it to be // parsed from YAML configuration files. It accepts both numeric values // (interpreted as bytes) and string values with units (e.g., "10MB"). func (s *Size) UnmarshalYAML(unmarshal func(interface{}) error) error { // Try to unmarshal as int64 first var intVal int64 if err := unmarshal(&intVal); err == nil { *s = Size(intVal) return nil } // Try to unmarshal as string var strVal string if err := unmarshal(&strVal); err != nil { return fmt.Errorf("size must be a number or string") } // Parse the string using go-humanize bytes, err := humanize.ParseBytes(strVal) if err != nil { return fmt.Errorf("invalid size format: %w", err) } *s = Size(bytes) return nil } // Int64 returns the size as int64 bytes. // This is useful when the size needs to be passed to APIs that expect // a numeric byte count. func (s Size) Int64() int64 { return int64(s) } // String returns the size as a human-readable string. // For example, 1048576 bytes would be formatted as "1.0 MB". // This implements the fmt.Stringer interface. func (s Size) String() string { return humanize.Bytes(uint64(s)) } // ParseSize parses a size string into a Size value func ParseSize(s string) (Size, error) { bytes, err := humanize.ParseBytes(s) if err != nil { return 0, fmt.Errorf("invalid size format: %w", err) } return Size(bytes), nil }