env.go 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276
  1. package engine
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "strconv"
  8. "strings"
  9. )
  10. type Env []string
  11. func (env *Env) Get(key string) (value string) {
  12. // FIXME: use Map()
  13. for _, kv := range *env {
  14. if strings.Index(kv, "=") == -1 {
  15. continue
  16. }
  17. parts := strings.SplitN(kv, "=", 2)
  18. if parts[0] != key {
  19. continue
  20. }
  21. if len(parts) < 2 {
  22. value = ""
  23. } else {
  24. value = parts[1]
  25. }
  26. }
  27. return
  28. }
  29. func (env *Env) Exists(key string) bool {
  30. _, exists := env.Map()[key]
  31. return exists
  32. }
  33. // Len returns the number of keys in the environment.
  34. // Note that len(env) might be different from env.Len(),
  35. // because the same key might be set multiple times.
  36. func (env *Env) Len() int {
  37. return len(env.Map())
  38. }
  39. func (env *Env) Init(src *Env) {
  40. (*env) = make([]string, 0, len(*src))
  41. for _, val := range *src {
  42. (*env) = append((*env), val)
  43. }
  44. }
  45. func (env *Env) GetBool(key string) (value bool) {
  46. s := strings.ToLower(strings.Trim(env.Get(key), " \t"))
  47. if s == "" || s == "0" || s == "no" || s == "false" || s == "none" {
  48. return false
  49. }
  50. return true
  51. }
  52. func (env *Env) SetBool(key string, value bool) {
  53. if value {
  54. env.Set(key, "1")
  55. } else {
  56. env.Set(key, "0")
  57. }
  58. }
  59. func (env *Env) GetInt(key string) int {
  60. return int(env.GetInt64(key))
  61. }
  62. func (env *Env) GetInt64(key string) int64 {
  63. s := strings.Trim(env.Get(key), " \t")
  64. val, err := strconv.ParseInt(s, 10, 64)
  65. if err != nil {
  66. return 0
  67. }
  68. return val
  69. }
  70. func (env *Env) SetInt(key string, value int) {
  71. env.Set(key, fmt.Sprintf("%d", value))
  72. }
  73. func (env *Env) SetInt64(key string, value int64) {
  74. env.Set(key, fmt.Sprintf("%d", value))
  75. }
  76. // Returns nil if key not found
  77. func (env *Env) GetList(key string) []string {
  78. sval := env.Get(key)
  79. if sval == "" {
  80. return nil
  81. }
  82. l := make([]string, 0, 1)
  83. if err := json.Unmarshal([]byte(sval), &l); err != nil {
  84. l = append(l, sval)
  85. }
  86. return l
  87. }
  88. func (env *Env) GetSubEnv(key string) *Env {
  89. sval := env.Get(key)
  90. if sval == "" {
  91. return nil
  92. }
  93. buf := bytes.NewBufferString(sval)
  94. var sub Env
  95. if err := sub.Decode(buf); err != nil {
  96. return nil
  97. }
  98. return &sub
  99. }
  100. func (env *Env) SetSubEnv(key string, sub *Env) error {
  101. var buf bytes.Buffer
  102. if err := sub.Encode(&buf); err != nil {
  103. return err
  104. }
  105. env.Set(key, string(buf.Bytes()))
  106. return nil
  107. }
  108. func (env *Env) GetJson(key string, iface interface{}) error {
  109. sval := env.Get(key)
  110. if sval == "" {
  111. return nil
  112. }
  113. return json.Unmarshal([]byte(sval), iface)
  114. }
  115. func (env *Env) SetJson(key string, value interface{}) error {
  116. sval, err := json.Marshal(value)
  117. if err != nil {
  118. return err
  119. }
  120. env.Set(key, string(sval))
  121. return nil
  122. }
  123. func (env *Env) SetList(key string, value []string) error {
  124. return env.SetJson(key, value)
  125. }
  126. func (env *Env) Set(key, value string) {
  127. *env = append(*env, key+"="+value)
  128. }
  129. func NewDecoder(src io.Reader) *Decoder {
  130. return &Decoder{
  131. json.NewDecoder(src),
  132. }
  133. }
  134. type Decoder struct {
  135. *json.Decoder
  136. }
  137. func (decoder *Decoder) Decode() (*Env, error) {
  138. m := make(map[string]interface{})
  139. if err := decoder.Decoder.Decode(&m); err != nil {
  140. return nil, err
  141. }
  142. env := &Env{}
  143. for key, value := range m {
  144. env.SetAuto(key, value)
  145. }
  146. return env, nil
  147. }
  148. // DecodeEnv decodes `src` as a json dictionary, and adds
  149. // each decoded key-value pair to the environment.
  150. //
  151. // If `src` cannot be decoded as a json dictionary, an error
  152. // is returned.
  153. func (env *Env) Decode(src io.Reader) error {
  154. m := make(map[string]interface{})
  155. if err := json.NewDecoder(src).Decode(&m); err != nil {
  156. return err
  157. }
  158. for k, v := range m {
  159. env.SetAuto(k, v)
  160. }
  161. return nil
  162. }
  163. func (env *Env) SetAuto(k string, v interface{}) {
  164. // FIXME: we fix-convert float values to int, because
  165. // encoding/json decodes integers to float64, but cannot encode them back.
  166. // (See http://golang.org/src/pkg/encoding/json/decode.go#L46)
  167. if fval, ok := v.(float64); ok {
  168. env.SetInt64(k, int64(fval))
  169. } else if sval, ok := v.(string); ok {
  170. env.Set(k, sval)
  171. } else if val, err := json.Marshal(v); err == nil {
  172. env.Set(k, string(val))
  173. } else {
  174. env.Set(k, fmt.Sprintf("%v", v))
  175. }
  176. }
  177. func (env *Env) Encode(dst io.Writer) error {
  178. m := make(map[string]interface{})
  179. for k, v := range env.Map() {
  180. var val interface{}
  181. if err := json.Unmarshal([]byte(v), &val); err == nil {
  182. // FIXME: we fix-convert float values to int, because
  183. // encoding/json decodes integers to float64, but cannot encode them back.
  184. // (See http://golang.org/src/pkg/encoding/json/decode.go#L46)
  185. if fval, isFloat := val.(float64); isFloat {
  186. val = int(fval)
  187. }
  188. m[k] = val
  189. } else {
  190. m[k] = v
  191. }
  192. }
  193. if err := json.NewEncoder(dst).Encode(&m); err != nil {
  194. return err
  195. }
  196. return nil
  197. }
  198. func (env *Env) WriteTo(dst io.Writer) (n int64, err error) {
  199. // FIXME: return the number of bytes written to respect io.WriterTo
  200. return 0, env.Encode(dst)
  201. }
  202. func (env *Env) Import(src interface{}) (err error) {
  203. defer func() {
  204. if err != nil {
  205. err = fmt.Errorf("ImportEnv: %s", err)
  206. }
  207. }()
  208. var buf bytes.Buffer
  209. if err := json.NewEncoder(&buf).Encode(src); err != nil {
  210. return err
  211. }
  212. if err := env.Decode(&buf); err != nil {
  213. return err
  214. }
  215. return nil
  216. }
  217. func (env *Env) Map() map[string]string {
  218. m := make(map[string]string)
  219. for _, kv := range *env {
  220. parts := strings.SplitN(kv, "=", 2)
  221. m[parts[0]] = parts[1]
  222. }
  223. return m
  224. }
  225. // MultiMap returns a representation of env as a
  226. // map of string arrays, keyed by string.
  227. // This is the same structure as http headers for example,
  228. // which allow each key to have multiple values.
  229. func (env *Env) MultiMap() map[string][]string {
  230. m := make(map[string][]string)
  231. for _, kv := range *env {
  232. parts := strings.SplitN(kv, "=", 2)
  233. m[parts[0]] = append(m[parts[0]], parts[1])
  234. }
  235. return m
  236. }
  237. // InitMultiMap removes all values in env, then initializes
  238. // new values from the contents of m.
  239. func (env *Env) InitMultiMap(m map[string][]string) {
  240. (*env) = make([]string, 0, len(m))
  241. for k, vals := range m {
  242. for _, v := range vals {
  243. env.Set(k, v)
  244. }
  245. }
  246. }