encode.go 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398
  1. package toml
  2. import (
  3. "bytes"
  4. "encoding"
  5. "fmt"
  6. "io"
  7. "reflect"
  8. "sort"
  9. "strconv"
  10. "time"
  11. "github.com/naoina/toml/ast"
  12. )
  13. const (
  14. tagOmitempty = "omitempty"
  15. tagSkip = "-"
  16. )
  17. // Marshal returns the TOML encoding of v.
  18. //
  19. // Struct values encode as TOML. Each exported struct field becomes a field of
  20. // the TOML structure unless
  21. // - the field's tag is "-", or
  22. // - the field is empty and its tag specifies the "omitempty" option.
  23. //
  24. // The "toml" key in the struct field's tag value is the key name, followed by
  25. // an optional comma and options. Examples:
  26. //
  27. // // Field is ignored by this package.
  28. // Field int `toml:"-"`
  29. //
  30. // // Field appears in TOML as key "myName".
  31. // Field int `toml:"myName"`
  32. //
  33. // // Field appears in TOML as key "myName" and the field is omitted from the
  34. // // result of encoding if its value is empty.
  35. // Field int `toml:"myName,omitempty"`
  36. //
  37. // // Field appears in TOML as key "field", but the field is skipped if
  38. // // empty. Note the leading comma.
  39. // Field int `toml:",omitempty"`
  40. func (cfg *Config) Marshal(v interface{}) ([]byte, error) {
  41. buf := new(bytes.Buffer)
  42. err := cfg.NewEncoder(buf).Encode(v)
  43. return buf.Bytes(), err
  44. }
  45. // A Encoder writes TOML to an output stream.
  46. type Encoder struct {
  47. w io.Writer
  48. cfg *Config
  49. }
  50. // NewEncoder returns a new Encoder that writes to w.
  51. func (cfg *Config) NewEncoder(w io.Writer) *Encoder {
  52. return &Encoder{w, cfg}
  53. }
  54. // Encode writes the TOML of v to the stream.
  55. // See the documentation for Marshal for details about the conversion of Go values to TOML.
  56. func (e *Encoder) Encode(v interface{}) error {
  57. rv := reflect.ValueOf(v)
  58. for rv.Kind() == reflect.Ptr {
  59. if rv.IsNil() {
  60. return &marshalNilError{rv.Type()}
  61. }
  62. rv = rv.Elem()
  63. }
  64. buf := &tableBuf{typ: ast.TableTypeNormal}
  65. var err error
  66. switch rv.Kind() {
  67. case reflect.Struct:
  68. err = buf.structFields(e.cfg, rv)
  69. case reflect.Map:
  70. err = buf.mapFields(e.cfg, rv)
  71. default:
  72. err = &marshalTableError{rv.Type()}
  73. }
  74. if err != nil {
  75. return err
  76. }
  77. return buf.writeTo(e.w, "")
  78. }
  79. // Marshaler can be implemented to override the encoding of TOML values. The returned text
  80. // must be a simple TOML value (i.e. not a table) and is inserted into marshaler output.
  81. //
  82. // This interface exists for backwards-compatibility reasons. You probably want to
  83. // implement encoding.TextMarshaler or MarshalerRec instead.
  84. type Marshaler interface {
  85. MarshalTOML() ([]byte, error)
  86. }
  87. // MarshalerRec can be implemented to override the TOML encoding of a type.
  88. // The returned value is marshaled in place of the receiver.
  89. type MarshalerRec interface {
  90. MarshalTOML() (interface{}, error)
  91. }
  92. type tableBuf struct {
  93. name string // already escaped / quoted
  94. body []byte
  95. children []*tableBuf
  96. typ ast.TableType
  97. arrayDepth int
  98. }
  99. func (b *tableBuf) writeTo(w io.Writer, prefix string) error {
  100. key := b.name // TODO: escape dots
  101. if prefix != "" {
  102. key = prefix + "." + key
  103. }
  104. if b.name != "" {
  105. head := "[" + key + "]"
  106. if b.typ == ast.TableTypeArray {
  107. head = "[" + head + "]"
  108. }
  109. head += "\n"
  110. if _, err := io.WriteString(w, head); err != nil {
  111. return err
  112. }
  113. }
  114. if _, err := w.Write(b.body); err != nil {
  115. return err
  116. }
  117. for i, child := range b.children {
  118. if len(b.body) > 0 || i > 0 {
  119. if _, err := w.Write([]byte("\n")); err != nil {
  120. return err
  121. }
  122. }
  123. if err := child.writeTo(w, key); err != nil {
  124. return err
  125. }
  126. }
  127. return nil
  128. }
  129. func (b *tableBuf) newChild(name string) *tableBuf {
  130. child := &tableBuf{name: quoteName(name), typ: ast.TableTypeNormal}
  131. if b.arrayDepth > 0 {
  132. child.typ = ast.TableTypeArray
  133. }
  134. return child
  135. }
  136. func (b *tableBuf) addChild(child *tableBuf) {
  137. // Empty table elision: we can avoid writing a table that doesn't have any keys on its
  138. // own. Array tables can't be elided because they define array elements (which would
  139. // be missing if elided).
  140. if len(child.body) == 0 && child.typ == ast.TableTypeNormal {
  141. for _, gchild := range child.children {
  142. gchild.name = child.name + "." + gchild.name
  143. b.addChild(gchild)
  144. }
  145. return
  146. }
  147. b.children = append(b.children, child)
  148. }
  149. func (b *tableBuf) structFields(cfg *Config, rv reflect.Value) error {
  150. rt := rv.Type()
  151. for i := 0; i < rv.NumField(); i++ {
  152. ft := rt.Field(i)
  153. if ft.PkgPath != "" && !ft.Anonymous { // not exported
  154. continue
  155. }
  156. name, rest := extractTag(ft.Tag.Get(fieldTagName))
  157. if name == tagSkip {
  158. continue
  159. }
  160. fv := rv.Field(i)
  161. if rest == tagOmitempty && isEmptyValue(fv) {
  162. continue
  163. }
  164. if name == "" {
  165. name = cfg.FieldToKey(rt, ft.Name)
  166. }
  167. if err := b.field(cfg, name, fv); err != nil {
  168. return err
  169. }
  170. }
  171. return nil
  172. }
  173. type mapKeyList []struct {
  174. key string
  175. value reflect.Value
  176. }
  177. func (l mapKeyList) Len() int { return len(l) }
  178. func (l mapKeyList) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
  179. func (l mapKeyList) Less(i, j int) bool { return l[i].key < l[j].key }
  180. func (b *tableBuf) mapFields(cfg *Config, rv reflect.Value) error {
  181. keys := rv.MapKeys()
  182. keylist := make(mapKeyList, len(keys))
  183. for i, key := range keys {
  184. var err error
  185. keylist[i].key, err = encodeMapKey(key)
  186. if err != nil {
  187. return err
  188. }
  189. keylist[i].value = rv.MapIndex(key)
  190. }
  191. sort.Sort(keylist)
  192. for _, kv := range keylist {
  193. if err := b.field(cfg, kv.key, kv.value); err != nil {
  194. return err
  195. }
  196. }
  197. return nil
  198. }
  199. func (b *tableBuf) field(cfg *Config, name string, rv reflect.Value) error {
  200. off := len(b.body)
  201. b.body = append(b.body, quoteName(name)...)
  202. b.body = append(b.body, " = "...)
  203. isTable, err := b.value(cfg, rv, name)
  204. if isTable {
  205. b.body = b.body[:off] // rub out "key ="
  206. } else {
  207. b.body = append(b.body, '\n')
  208. }
  209. return err
  210. }
  211. func (b *tableBuf) value(cfg *Config, rv reflect.Value, name string) (bool, error) {
  212. isMarshaler, isTable, err := b.marshaler(cfg, rv, name)
  213. if isMarshaler {
  214. return isTable, err
  215. }
  216. switch rv.Kind() {
  217. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  218. b.body = strconv.AppendInt(b.body, rv.Int(), 10)
  219. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
  220. b.body = strconv.AppendUint(b.body, rv.Uint(), 10)
  221. case reflect.Float32, reflect.Float64:
  222. b.body = strconv.AppendFloat(b.body, rv.Float(), 'e', -1, 64)
  223. case reflect.Bool:
  224. b.body = strconv.AppendBool(b.body, rv.Bool())
  225. case reflect.String:
  226. b.body = strconv.AppendQuote(b.body, rv.String())
  227. case reflect.Ptr, reflect.Interface:
  228. if rv.IsNil() {
  229. return false, &marshalNilError{rv.Type()}
  230. }
  231. return b.value(cfg, rv.Elem(), name)
  232. case reflect.Slice, reflect.Array:
  233. rvlen := rv.Len()
  234. if rvlen == 0 {
  235. b.body = append(b.body, '[', ']')
  236. return false, nil
  237. }
  238. b.arrayDepth++
  239. wroteElem := false
  240. b.body = append(b.body, '[')
  241. for i := 0; i < rvlen; i++ {
  242. isTable, err := b.value(cfg, rv.Index(i), name)
  243. if err != nil {
  244. return isTable, err
  245. }
  246. wroteElem = wroteElem || !isTable
  247. if wroteElem {
  248. if i < rvlen-1 {
  249. b.body = append(b.body, ',', ' ')
  250. } else {
  251. b.body = append(b.body, ']')
  252. }
  253. }
  254. }
  255. if !wroteElem {
  256. b.body = b.body[:len(b.body)-1] // rub out '['
  257. }
  258. b.arrayDepth--
  259. return !wroteElem, nil
  260. case reflect.Struct:
  261. child := b.newChild(name)
  262. err := child.structFields(cfg, rv)
  263. b.addChild(child)
  264. return true, err
  265. case reflect.Map:
  266. child := b.newChild(name)
  267. err := child.mapFields(cfg, rv)
  268. b.addChild(child)
  269. return true, err
  270. default:
  271. return false, fmt.Errorf("toml: marshal: unsupported type %v", rv.Kind())
  272. }
  273. return false, nil
  274. }
  275. func (b *tableBuf) marshaler(cfg *Config, rv reflect.Value, name string) (handled, isTable bool, err error) {
  276. switch t := rv.Interface().(type) {
  277. case encoding.TextMarshaler:
  278. enc, err := t.MarshalText()
  279. if err != nil {
  280. return true, false, err
  281. }
  282. b.body = encodeTextMarshaler(b.body, string(enc))
  283. return true, false, nil
  284. case MarshalerRec:
  285. newval, err := t.MarshalTOML()
  286. if err != nil {
  287. return true, false, err
  288. }
  289. isTable, err = b.value(cfg, reflect.ValueOf(newval), name)
  290. return true, isTable, err
  291. case Marshaler:
  292. enc, err := t.MarshalTOML()
  293. if err != nil {
  294. return true, false, err
  295. }
  296. b.body = append(b.body, enc...)
  297. return true, false, nil
  298. }
  299. return false, false, nil
  300. }
  301. func encodeTextMarshaler(buf []byte, v string) []byte {
  302. // Emit the value without quotes if possible.
  303. if v == "true" || v == "false" {
  304. return append(buf, v...)
  305. } else if _, err := time.Parse(time.RFC3339Nano, v); err == nil {
  306. return append(buf, v...)
  307. } else if _, err := strconv.ParseInt(v, 10, 64); err == nil {
  308. return append(buf, v...)
  309. } else if _, err := strconv.ParseUint(v, 10, 64); err == nil {
  310. return append(buf, v...)
  311. } else if _, err := strconv.ParseFloat(v, 64); err == nil {
  312. return append(buf, v...)
  313. }
  314. return strconv.AppendQuote(buf, v)
  315. }
  316. func encodeMapKey(rv reflect.Value) (string, error) {
  317. if rv.Kind() == reflect.String {
  318. return rv.String(), nil
  319. }
  320. if tm, ok := rv.Interface().(encoding.TextMarshaler); ok {
  321. b, err := tm.MarshalText()
  322. return string(b), err
  323. }
  324. switch rv.Kind() {
  325. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  326. return strconv.FormatInt(rv.Int(), 10), nil
  327. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
  328. return strconv.FormatUint(rv.Uint(), 10), nil
  329. }
  330. return "", fmt.Errorf("toml: invalid map key type %v", rv.Type())
  331. }
  332. func isEmptyValue(v reflect.Value) bool {
  333. switch v.Kind() {
  334. case reflect.Array:
  335. // encoding/json treats all arrays with non-zero length as non-empty. We check the
  336. // array content here because zero-length arrays are almost never used.
  337. len := v.Len()
  338. for i := 0; i < len; i++ {
  339. if !isEmptyValue(v.Index(i)) {
  340. return false
  341. }
  342. }
  343. return true
  344. case reflect.Map, reflect.Slice, reflect.String:
  345. return v.Len() == 0
  346. case reflect.Bool:
  347. return !v.Bool()
  348. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  349. return v.Int() == 0
  350. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
  351. return v.Uint() == 0
  352. case reflect.Float32, reflect.Float64:
  353. return v.Float() == 0
  354. case reflect.Interface, reflect.Ptr:
  355. return v.IsNil()
  356. }
  357. return false
  358. }
  359. func quoteName(s string) string {
  360. if len(s) == 0 {
  361. return strconv.Quote(s)
  362. }
  363. for _, r := range s {
  364. if r >= '0' && r <= '9' || r >= 'A' && r <= 'Z' || r >= 'a' && r <= 'z' || r == '-' || r == '_' {
  365. continue
  366. }
  367. return strconv.Quote(s)
  368. }
  369. return s
  370. }