env.go 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385
  1. package engine
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "sort"
  8. "strconv"
  9. "strings"
  10. )
  11. type Env []string
  12. func (env *Env) Get(key string) (value string) {
  13. // FIXME: use Map()
  14. for _, kv := range *env {
  15. if strings.Index(kv, "=") == -1 {
  16. continue
  17. }
  18. parts := strings.SplitN(kv, "=", 2)
  19. if parts[0] != key {
  20. continue
  21. }
  22. if len(parts) < 2 {
  23. value = ""
  24. } else {
  25. value = parts[1]
  26. }
  27. }
  28. return
  29. }
  30. func (env *Env) Exists(key string) bool {
  31. _, exists := env.Map()[key]
  32. return exists
  33. }
  34. // Len returns the number of keys in the environment.
  35. // Note that len(env) might be different from env.Len(),
  36. // because the same key might be set multiple times.
  37. func (env *Env) Len() int {
  38. return len(env.Map())
  39. }
  40. func (env *Env) Init(src *Env) {
  41. (*env) = make([]string, 0, len(*src))
  42. for _, val := range *src {
  43. (*env) = append((*env), val)
  44. }
  45. }
  46. func (env *Env) GetBool(key string) (value bool) {
  47. s := strings.ToLower(strings.Trim(env.Get(key), " \t"))
  48. if s == "" || s == "0" || s == "no" || s == "false" || s == "none" {
  49. return false
  50. }
  51. return true
  52. }
  53. func (env *Env) SetBool(key string, value bool) {
  54. if value {
  55. env.Set(key, "1")
  56. } else {
  57. env.Set(key, "0")
  58. }
  59. }
  60. func (env *Env) GetInt(key string) int {
  61. return int(env.GetInt64(key))
  62. }
  63. func (env *Env) GetInt64(key string) int64 {
  64. s := strings.Trim(env.Get(key), " \t")
  65. val, err := strconv.ParseInt(s, 10, 64)
  66. if err != nil {
  67. return 0
  68. }
  69. return val
  70. }
  71. func (env *Env) SetInt(key string, value int) {
  72. env.Set(key, fmt.Sprintf("%d", value))
  73. }
  74. func (env *Env) SetInt64(key string, value int64) {
  75. env.Set(key, fmt.Sprintf("%d", value))
  76. }
  77. // Returns nil if key not found
  78. func (env *Env) GetList(key string) []string {
  79. sval := env.Get(key)
  80. if sval == "" {
  81. return nil
  82. }
  83. l := make([]string, 0, 1)
  84. if err := json.Unmarshal([]byte(sval), &l); err != nil {
  85. l = append(l, sval)
  86. }
  87. return l
  88. }
  89. func (env *Env) GetSubEnv(key string) *Env {
  90. sval := env.Get(key)
  91. if sval == "" {
  92. return nil
  93. }
  94. buf := bytes.NewBufferString(sval)
  95. var sub Env
  96. if err := sub.Decode(buf); err != nil {
  97. return nil
  98. }
  99. return &sub
  100. }
  101. func (env *Env) SetSubEnv(key string, sub *Env) error {
  102. var buf bytes.Buffer
  103. if err := sub.Encode(&buf); err != nil {
  104. return err
  105. }
  106. env.Set(key, string(buf.Bytes()))
  107. return nil
  108. }
  109. func (env *Env) GetJson(key string, iface interface{}) error {
  110. sval := env.Get(key)
  111. if sval == "" {
  112. return nil
  113. }
  114. return json.Unmarshal([]byte(sval), iface)
  115. }
  116. func (env *Env) SetJson(key string, value interface{}) error {
  117. sval, err := json.Marshal(value)
  118. if err != nil {
  119. return err
  120. }
  121. env.Set(key, string(sval))
  122. return nil
  123. }
  124. func (env *Env) SetList(key string, value []string) error {
  125. return env.SetJson(key, value)
  126. }
  127. func (env *Env) Set(key, value string) {
  128. *env = append(*env, key+"="+value)
  129. }
  130. func NewDecoder(src io.Reader) *Decoder {
  131. return &Decoder{
  132. json.NewDecoder(src),
  133. }
  134. }
  135. type Decoder struct {
  136. *json.Decoder
  137. }
  138. func (decoder *Decoder) Decode() (*Env, error) {
  139. m := make(map[string]interface{})
  140. if err := decoder.Decoder.Decode(&m); err != nil {
  141. return nil, err
  142. }
  143. env := &Env{}
  144. for key, value := range m {
  145. env.SetAuto(key, value)
  146. }
  147. return env, nil
  148. }
  149. // DecodeEnv decodes `src` as a json dictionary, and adds
  150. // each decoded key-value pair to the environment.
  151. //
  152. // If `src` cannot be decoded as a json dictionary, an error
  153. // is returned.
  154. func (env *Env) Decode(src io.Reader) error {
  155. m := make(map[string]interface{})
  156. if err := json.NewDecoder(src).Decode(&m); err != nil {
  157. return err
  158. }
  159. for k, v := range m {
  160. env.SetAuto(k, v)
  161. }
  162. return nil
  163. }
  164. func (env *Env) SetAuto(k string, v interface{}) {
  165. // FIXME: we fix-convert float values to int, because
  166. // encoding/json decodes integers to float64, but cannot encode them back.
  167. // (See http://golang.org/src/pkg/encoding/json/decode.go#L46)
  168. if fval, ok := v.(float64); ok {
  169. env.SetInt64(k, int64(fval))
  170. } else if sval, ok := v.(string); ok {
  171. env.Set(k, sval)
  172. } else if val, err := json.Marshal(v); err == nil {
  173. env.Set(k, string(val))
  174. } else {
  175. env.Set(k, fmt.Sprintf("%v", v))
  176. }
  177. }
  178. func (env *Env) Encode(dst io.Writer) error {
  179. m := make(map[string]interface{})
  180. for k, v := range env.Map() {
  181. var val interface{}
  182. if err := json.Unmarshal([]byte(v), &val); err == nil {
  183. // FIXME: we fix-convert float values to int, because
  184. // encoding/json decodes integers to float64, but cannot encode them back.
  185. // (See http://golang.org/src/pkg/encoding/json/decode.go#L46)
  186. if fval, isFloat := val.(float64); isFloat {
  187. val = int(fval)
  188. }
  189. m[k] = val
  190. } else {
  191. m[k] = v
  192. }
  193. }
  194. if err := json.NewEncoder(dst).Encode(&m); err != nil {
  195. return err
  196. }
  197. return nil
  198. }
  199. func (env *Env) WriteTo(dst io.Writer) (n int64, err error) {
  200. // FIXME: return the number of bytes written to respect io.WriterTo
  201. return 0, env.Encode(dst)
  202. }
  203. func (env *Env) Import(src interface{}) (err error) {
  204. defer func() {
  205. if err != nil {
  206. err = fmt.Errorf("ImportEnv: %s", err)
  207. }
  208. }()
  209. var buf bytes.Buffer
  210. if err := json.NewEncoder(&buf).Encode(src); err != nil {
  211. return err
  212. }
  213. if err := env.Decode(&buf); err != nil {
  214. return err
  215. }
  216. return nil
  217. }
  218. func (env *Env) Map() map[string]string {
  219. m := make(map[string]string)
  220. for _, kv := range *env {
  221. parts := strings.SplitN(kv, "=", 2)
  222. m[parts[0]] = parts[1]
  223. }
  224. return m
  225. }
  226. type Table struct {
  227. Data []*Env
  228. sortKey string
  229. Chan chan *Env
  230. }
  231. func NewTable(sortKey string, sizeHint int) *Table {
  232. return &Table{
  233. make([]*Env, 0, sizeHint),
  234. sortKey,
  235. make(chan *Env),
  236. }
  237. }
  238. func (t *Table) SetKey(sortKey string) {
  239. t.sortKey = sortKey
  240. }
  241. func (t *Table) Add(env *Env) {
  242. t.Data = append(t.Data, env)
  243. }
  244. func (t *Table) Len() int {
  245. return len(t.Data)
  246. }
  247. func (t *Table) Less(a, b int) bool {
  248. return t.lessBy(a, b, t.sortKey)
  249. }
  250. func (t *Table) lessBy(a, b int, by string) bool {
  251. keyA := t.Data[a].Get(by)
  252. keyB := t.Data[b].Get(by)
  253. intA, errA := strconv.ParseInt(keyA, 10, 64)
  254. intB, errB := strconv.ParseInt(keyB, 10, 64)
  255. if errA == nil && errB == nil {
  256. return intA < intB
  257. }
  258. return keyA < keyB
  259. }
  260. func (t *Table) Swap(a, b int) {
  261. tmp := t.Data[a]
  262. t.Data[a] = t.Data[b]
  263. t.Data[b] = tmp
  264. }
  265. func (t *Table) Sort() {
  266. sort.Sort(t)
  267. }
  268. func (t *Table) ReverseSort() {
  269. sort.Sort(sort.Reverse(t))
  270. }
  271. func (t *Table) WriteListTo(dst io.Writer) (n int64, err error) {
  272. if _, err := dst.Write([]byte{'['}); err != nil {
  273. return -1, err
  274. }
  275. n = 1
  276. for i, env := range t.Data {
  277. bytes, err := env.WriteTo(dst)
  278. if err != nil {
  279. return -1, err
  280. }
  281. n += bytes
  282. if i != len(t.Data)-1 {
  283. if _, err := dst.Write([]byte{','}); err != nil {
  284. return -1, err
  285. }
  286. n += 1
  287. }
  288. }
  289. if _, err := dst.Write([]byte{']'}); err != nil {
  290. return -1, err
  291. }
  292. return n + 1, nil
  293. }
  294. func (t *Table) ToListString() (string, error) {
  295. buffer := bytes.NewBuffer(nil)
  296. if _, err := t.WriteListTo(buffer); err != nil {
  297. return "", err
  298. }
  299. return buffer.String(), nil
  300. }
  301. func (t *Table) WriteTo(dst io.Writer) (n int64, err error) {
  302. for _, env := range t.Data {
  303. bytes, err := env.WriteTo(dst)
  304. if err != nil {
  305. return -1, err
  306. }
  307. n += bytes
  308. }
  309. return n, nil
  310. }
  311. func (t *Table) ReadListFrom(src []byte) (n int64, err error) {
  312. var array []interface{}
  313. if err := json.Unmarshal(src, &array); err != nil {
  314. return -1, err
  315. }
  316. for _, item := range array {
  317. if m, ok := item.(map[string]interface{}); ok {
  318. env := &Env{}
  319. for key, value := range m {
  320. env.SetAuto(key, value)
  321. }
  322. t.Add(env)
  323. }
  324. }
  325. return int64(len(src)), nil
  326. }
  327. func (t *Table) ReadFrom(src io.Reader) (n int64, err error) {
  328. decoder := NewDecoder(src)
  329. for {
  330. env, err := decoder.Decode()
  331. if err == io.EOF {
  332. return 0, nil
  333. } else if err != nil {
  334. return -1, err
  335. }
  336. t.Add(env)
  337. }
  338. return 0, nil
  339. }