You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

9 年之前
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243
  1. package env
  2. import (
  3. "bufio"
  4. "fmt"
  5. "os"
  6. "path/filepath"
  7. "reflect"
  8. "regexp"
  9. "strconv"
  10. "strings"
  11. log "github.com/Sirupsen/logrus"
  12. "github.com/goanywhere/x/fs"
  13. )
  14. const tag string = "env"
  15. var pattern = regexp.MustCompile(`(export\s*)?(?P<key>\w+)\s*(=)\s*(?P<value>("|')?[\w\s-$,:/\.\+]+)("|')?`)
  16. // findKeyValue finds '=' separated key/value pair from the given string.
  17. func findKeyValue(line string) (key, value string) {
  18. // Intentionally insert a linebreak here to avoid non-stop characters on [[:graph:]].
  19. line = fmt.Sprintf("%s\n", line)
  20. match := pattern.FindStringSubmatch(line)
  21. if len(match) != 0 {
  22. result := make(map[string]string)
  23. for index, name := range pattern.SubexpNames() {
  24. result[name] = match[index]
  25. }
  26. // preserve those quoted spaces for value.
  27. key, value = result["key"], strings.TrimSpace(result["value"])
  28. value = strings.Trim(value, "'")
  29. value = strings.Trim(value, "\"")
  30. }
  31. return
  32. }
  33. // Load parses & set the values in the given files into os environment.
  34. func Load(dotenv string) {
  35. if fs.Exists(dotenv) {
  36. dotenv, _ = filepath.Abs(dotenv)
  37. if file, err := os.Open(dotenv); err == nil {
  38. defer file.Close()
  39. scanner := bufio.NewScanner(file)
  40. for scanner.Scan() {
  41. k, v := findKeyValue(scanner.Text())
  42. if k != "" && v != "" {
  43. Set(k, v)
  44. } else {
  45. continue
  46. }
  47. }
  48. }
  49. }
  50. }
  51. // Get retrieves the string value of the environment variable named by the key.
  52. // It returns the value, which will be empty if the variable is not present.
  53. func Get(key string) (value string, exists bool) {
  54. if v := os.Getenv(key); v != "" {
  55. value, exists = v, true
  56. }
  57. return
  58. }
  59. // Set stores the value of the environment variable named by the key. It returns an error, if any.
  60. func Set(key string, value interface{}) error {
  61. var sv string
  62. switch T := value.(type) {
  63. case bool:
  64. sv = strconv.FormatBool(T)
  65. case float32, float64:
  66. sv = strconv.FormatFloat(reflect.ValueOf(value).Float(), 'g', -1, 64)
  67. case int, int8, int16, int32, int64:
  68. sv = strconv.FormatInt(reflect.ValueOf(value).Int(), 10)
  69. case uint, uint8, uint16, uint32, uint64:
  70. sv = strconv.FormatUint(reflect.ValueOf(value).Uint(), 10)
  71. case string:
  72. sv = value.(string)
  73. case []string:
  74. sv = strings.Join(value.([]string), ",")
  75. default:
  76. return fmt.Errorf("Unsupported type: %v", T)
  77. }
  78. return os.Setenv(key, sv)
  79. }
  80. // String retrieves the string value from environment named by the key.
  81. func String(key string, fallback ...string) (value string) {
  82. if str, exists := Get(key); exists {
  83. value = str
  84. } else if len(fallback) > 0 {
  85. value = fallback[0]
  86. }
  87. return
  88. }
  89. // Strings retrieves the string values separated by comma from the environment.
  90. func Strings(key string, fallback ...[]string) (value []string) {
  91. if v, exists := Get(key); exists {
  92. for _, item := range strings.Split(strings.TrimSpace(v), ",") {
  93. value = append(value, strings.TrimSpace(item))
  94. }
  95. } else if len(fallback) > 0 {
  96. value = fallback[0]
  97. }
  98. return
  99. }
  100. // Int retrieves the integer values separated by comma from the environment.
  101. func Int(key string, fallback ...int) (value int) {
  102. if str, exists := Get(key); exists {
  103. if v, e := strconv.ParseInt(str, 10, 0); e == nil {
  104. value = int(v)
  105. }
  106. } else if len(fallback) > 0 {
  107. value = fallback[0]
  108. }
  109. return
  110. }
  111. // Int retrieves the 64-bit integer values separated by comma from the environment.
  112. func Int64(key string, fallback ...int64) (value int64) {
  113. if str, exists := Get(key); exists {
  114. if v, e := strconv.ParseInt(str, 10, 64); e == nil {
  115. value = v
  116. }
  117. } else if len(fallback) > 0 {
  118. value = fallback[0]
  119. }
  120. return
  121. }
  122. // Uint retrieves the unsigned integer values separated by comma from the environment.
  123. func Uint(key string, fallback ...uint) (value uint) {
  124. if str, exists := Get(key); exists {
  125. if v, e := strconv.ParseUint(str, 10, 0); e == nil {
  126. value = uint(v)
  127. }
  128. } else if len(fallback) > 0 {
  129. value = fallback[0]
  130. }
  131. return
  132. }
  133. // Uint64 retrieves the 64-bit unsigned integer values separated by comma from the environment.
  134. func Uint64(key string, fallback ...uint64) (value uint64) {
  135. if str, exists := Get(key); exists {
  136. if v, e := strconv.ParseUint(str, 10, 64); e == nil {
  137. value = v
  138. }
  139. } else if len(fallback) > 0 {
  140. value = fallback[0]
  141. }
  142. return
  143. }
  144. // Bool retrieves boolean value from the environment.
  145. func Bool(key string, fallback ...bool) (value bool) {
  146. if str, exists := Get(key); exists {
  147. if v, e := strconv.ParseBool(str); e == nil {
  148. value = v
  149. }
  150. } else if len(fallback) > 0 {
  151. value = fallback[0]
  152. }
  153. return
  154. }
  155. // Float retrieves float (64) value from the environment.
  156. func Float(key string, fallback ...float64) (value float64) {
  157. if str, exists := Get(key); exists {
  158. if v, e := strconv.ParseFloat(str, 64); e == nil {
  159. value = v
  160. }
  161. } else if len(fallback) > 0 {
  162. value = fallback[0]
  163. }
  164. return
  165. }
  166. // Map fetches the key/value pair from os.Environ into the given spec.
  167. // Tags are supported via `env:ACTUAL_OS_KEY`.
  168. func Map(spec interface{}) error {
  169. value := reflect.ValueOf(spec)
  170. s := value.Elem()
  171. var stype reflect.Type = s.Type()
  172. var field reflect.Value
  173. for index := 0; index < s.NumField(); index++ {
  174. field = s.Field(index)
  175. if field.CanSet() {
  176. key := stype.Field(index).Tag.Get(tag)
  177. if key == "" {
  178. key = stype.Field(index).Name
  179. }
  180. value, exists := Get(key)
  181. if !exists {
  182. continue
  183. }
  184. // converts the environmental value from string to its real type.
  185. // Supports: String | Strings | Bool | Float | Integer | Unsiged Integer
  186. switch field.Kind() {
  187. case reflect.String:
  188. field.SetString(value)
  189. case reflect.Bool:
  190. field.SetBool(Bool(key))
  191. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  192. field.SetInt(Int64(key))
  193. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
  194. field.SetUint(Uint64(key))
  195. case reflect.Float32, reflect.Float64:
  196. field.SetFloat(Float(key))
  197. case reflect.Slice:
  198. switch field.Interface().(type) {
  199. case []string:
  200. field.Set(reflect.ValueOf(Strings(key)))
  201. default:
  202. log.Fatalf("Only string slice is supported")
  203. }
  204. }
  205. }
  206. }
  207. return nil
  208. }