sorbet/pkg/env/environ.go

package env

import (
	"errors"
	"fmt"
	"github.com/joho/godotenv"
	"reflect"
	"sorbet/pkg/cast"
	"strconv"
	"strings"
	"time"
	"unsafe"
)

type Environ map[string]string

func New() Environ {
	return make(Environ)
}

// Load 加载环境变量文件
func (e Environ) Load(filenames ...string) error {
	data, err := godotenv.Read(filenames...)
	if err != nil {
		return err
	}
	for key, value := range data {
		e[key] = value
	}
	return nil
}

// Lookup 查看环境变量值，如果不存在或值为空，返回的第二个参数的值则为false。
func (e Environ) Lookup(key string) (string, bool) {
	val, exists := e[key]
	if exists && len(val) == 0 {
		exists = false
	}
	return val, exists
}

// Exists 判断环境变量是否存在
func (e Environ) Exists(key string) bool {
	_, exists := e[key]
	return exists
}

// String 取字符串值
func (e Environ) String(key string, fallback ...string) string {
	if value, exists := e.Lookup(key); exists {
		return value
	}
	for _, value := range fallback {
		return value
	}
	return ""
}

// Bytes 取二进制值
func (e Environ) Bytes(key string, fallback ...[]byte) []byte {
	if value, exists := e.Lookup(key); exists {
		return []byte(value)
	}
	for _, bytes := range fallback {
		return bytes
	}
	return []byte{}
}

// Int 取整型值
func (e Environ) Int(key string, fallback ...int) int {
	if val, exists := e.Lookup(key); exists {
		if n, err := strconv.Atoi(val); err == nil {
			return n
		}
	}
	for _, value := range fallback {
		return value
	}
	return 0
}

func (e Environ) Duration(key string, fallback ...time.Duration) time.Duration {
	if val, ok := e.Lookup(key); ok {
		n, err := strconv.Atoi(val)
		if err == nil {
			return time.Duration(n)
		}
		if d, err := time.ParseDuration(val); err == nil {
			return d
		}
	}
	for _, value := range fallback {
		return value
	}
	return time.Duration(0)
}

func (e Environ) Bool(key string, fallback ...bool) bool {
	if val, ok := e.Lookup(key); ok {
		bl, err := strconv.ParseBool(val)
		if err == nil {
			return bl
		}
	}
	for _, value := range fallback {
		return value
	}
	return false
}

// List 将值按 `,` 分割并返回
func (e Environ) List(key string, fallback ...[]string) []string {
	if value, ok := e.Lookup(key); ok {
		parts := strings.Split(value, ",")
		for i, part := range parts {
			parts[i] = strings.Trim(part, " \n\r")
		}
		return parts
	}
	for _, value := range fallback {
		return value
	}
	return []string{}
}

// Map 获取指定前缀的所有值
func (e Environ) Map(prefix string) map[string]string {
	result := map[string]string{}
	for k, v := range e {
		if strings.HasPrefix(k, prefix) {
			name := strings.TrimPrefix(k, prefix)
			result[name] = strings.TrimSpace(v)
		}
	}
	return result
}

// Where 获取符合过滤器的所有值
func (e Environ) Where(filter func(name, value string) bool) map[string]string {
	result := map[string]string{}
	for k, v := range e {
		if filter(k, v) {
			result[k] = v
		}
	}
	return result
}

// Fill 将环境变量填充到指定结构体
func (e Environ) Fill(structure any) error {
	inputType := reflect.TypeOf(structure)

	if inputType != nil && inputType.Kind() == reflect.Ptr && inputType.Elem().Kind() == reflect.Struct {
		return e.fillStruct(reflect.ValueOf(structure).Elem())
	}

	return errors.New("env: invalid structure")
}

func (e Environ) fillStruct(s reflect.Value) error {
	for i := 0; i < s.NumField(); i++ {
		if t, exist := s.Type().Field(i).Tag.Lookup("env"); exist {
			if osv := e[t]; osv != "" {
				v, err := cast.FromType(osv, s.Type().Field(i).Type)
				if err != nil {
					return fmt.Errorf("env: cannot set `%v` field; err: %v", s.Type().Field(i).Name, err)
				}
				ptr := reflect.NewAt(s.Field(i).Type(), unsafe.Pointer(s.Field(i).UnsafeAddr())).Elem()
				ptr.Set(reflect.ValueOf(v))
			}
		} else if s.Type().Field(i).Type.Kind() == reflect.Struct {
			if err := e.fillStruct(s.Field(i)); err != nil {
				return err
			}
		} else if s.Type().Field(i).Type.Kind() == reflect.Ptr {
			if s.Field(i).IsZero() == false && s.Field(i).Elem().Type().Kind() == reflect.Struct {
				if err := e.fillStruct(s.Field(i).Elem()); err != nil {
					return err
				}
			}
		}
	}

	return nil
}