human_readable_bytes.go 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. package config
  2. import (
  3. "fmt"
  4. "reflect"
  5. "strings"
  6. "github.com/dustin/go-humanize"
  7. mapstructure "github.com/go-viper/mapstructure/v2"
  8. )
  9. // HumanReadableBytes is a byte size from config: a plain number (bytes) or a string
  10. // parsed with github.com/dustin/go-humanize (e.g. "10 MB", "1 GiB", "512 kb").
  11. type HumanReadableBytes int64
  12. // UnmarshalText implements encoding.TextUnmarshaler for YAML string values.
  13. func (h *HumanReadableBytes) UnmarshalText(text []byte) error {
  14. s := strings.TrimSpace(string(text))
  15. if s == "" {
  16. *h = 0
  17. return nil
  18. }
  19. n, err := humanize.ParseBytes(s)
  20. if err != nil {
  21. return fmt.Errorf("parse file size: %w", err)
  22. }
  23. *h = HumanReadableBytes(n)
  24. return nil
  25. }
  26. func humanReadableBytesNumericHook() mapstructure.DecodeHookFunc {
  27. var zero HumanReadableBytes
  28. target := reflect.TypeOf(zero)
  29. return func(from reflect.Type, to reflect.Type, data any) (any, error) {
  30. if to != target {
  31. return data, nil
  32. }
  33. if v, ok := humanReadableBytesFromDecodedValue(data); ok {
  34. return v, nil
  35. }
  36. return data, nil
  37. }
  38. }
  39. func humanReadableBytesFromDecodedValue(data any) (HumanReadableBytes, bool) {
  40. v, ok := ptrHumanReadableBytes(data)
  41. if ok {
  42. return v, true
  43. }
  44. if v, ok := data.(HumanReadableBytes); ok {
  45. return v, true
  46. }
  47. n, ok := int64FromYAMLNumber(data)
  48. if !ok {
  49. return 0, false
  50. }
  51. return HumanReadableBytes(n), true
  52. }
  53. func ptrHumanReadableBytes(data any) (HumanReadableBytes, bool) {
  54. p, ok := data.(*HumanReadableBytes)
  55. if !ok {
  56. return 0, false
  57. }
  58. if p == nil {
  59. return 0, true
  60. }
  61. return *p, true
  62. }
  63. func int64FromYAMLNumber(data any) (int64, bool) {
  64. if v, ok := int64FromFloatOrInt(data); ok {
  65. return v, true
  66. }
  67. if v, ok := data.(int64); ok {
  68. return v, true
  69. }
  70. if v, ok := data.(uint64); ok {
  71. return int64(v), true
  72. }
  73. return 0, false
  74. }
  75. func int64FromFloatOrInt(data any) (int64, bool) {
  76. if v, ok := data.(float64); ok {
  77. return int64(v), true
  78. }
  79. if v, ok := data.(int); ok {
  80. return int64(v), true
  81. }
  82. return 0, false
  83. }
  84. func newDefaultUnmarshalDecoderConfig() *mapstructure.DecoderConfig {
  85. return &mapstructure.DecoderConfig{
  86. DecodeHook: mapstructure.ComposeDecodeHookFunc(
  87. mapstructure.StringToTimeDurationHookFunc(),
  88. mapstructure.TextUnmarshallerHookFunc(),
  89. humanReadableBytesNumericHook(),
  90. ),
  91. WeaklyTypedInput: true,
  92. }
  93. }