encode.go 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551
  1. package toml
  2. import (
  3. "bufio"
  4. "errors"
  5. "fmt"
  6. "io"
  7. "reflect"
  8. "sort"
  9. "strconv"
  10. "strings"
  11. "time"
  12. )
  13. type tomlEncodeError struct{ error }
  14. var (
  15. errArrayMixedElementTypes = errors.New(
  16. "can't encode array with mixed element types")
  17. errArrayNilElement = errors.New(
  18. "can't encode array with nil element")
  19. errNonString = errors.New(
  20. "can't encode a map with non-string key type")
  21. errAnonNonStruct = errors.New(
  22. "can't encode an anonymous field that is not a struct")
  23. errArrayNoTable = errors.New(
  24. "TOML array element can't contain a table")
  25. errNoKey = errors.New(
  26. "top-level values must be a Go map or struct")
  27. errAnything = errors.New("") // used in testing
  28. )
  29. var quotedReplacer = strings.NewReplacer(
  30. "\t", "\\t",
  31. "\n", "\\n",
  32. "\r", "\\r",
  33. "\"", "\\\"",
  34. "\\", "\\\\",
  35. )
  36. // Encoder controls the encoding of Go values to a TOML document to some
  37. // io.Writer.
  38. //
  39. // The indentation level can be controlled with the Indent field.
  40. type Encoder struct {
  41. // A single indentation level. By default it is two spaces.
  42. Indent string
  43. // hasWritten is whether we have written any output to w yet.
  44. hasWritten bool
  45. w *bufio.Writer
  46. }
  47. // NewEncoder returns a TOML encoder that encodes Go values to the io.Writer
  48. // given. By default, a single indentation level is 2 spaces.
  49. func NewEncoder(w io.Writer) *Encoder {
  50. return &Encoder{
  51. w: bufio.NewWriter(w),
  52. Indent: " ",
  53. }
  54. }
  55. // Encode writes a TOML representation of the Go value to the underlying
  56. // io.Writer. If the value given cannot be encoded to a valid TOML document,
  57. // then an error is returned.
  58. //
  59. // The mapping between Go values and TOML values should be precisely the same
  60. // as for the Decode* functions. Similarly, the TextMarshaler interface is
  61. // supported by encoding the resulting bytes as strings. (If you want to write
  62. // arbitrary binary data then you will need to use something like base64 since
  63. // TOML does not have any binary types.)
  64. //
  65. // When encoding TOML hashes (i.e., Go maps or structs), keys without any
  66. // sub-hashes are encoded first.
  67. //
  68. // If a Go map is encoded, then its keys are sorted alphabetically for
  69. // deterministic output. More control over this behavior may be provided if
  70. // there is demand for it.
  71. //
  72. // Encoding Go values without a corresponding TOML representation---like map
  73. // types with non-string keys---will cause an error to be returned. Similarly
  74. // for mixed arrays/slices, arrays/slices with nil elements, embedded
  75. // non-struct types and nested slices containing maps or structs.
  76. // (e.g., [][]map[string]string is not allowed but []map[string]string is OK
  77. // and so is []map[string][]string.)
  78. func (enc *Encoder) Encode(v interface{}) error {
  79. rv := eindirect(reflect.ValueOf(v))
  80. if err := enc.safeEncode(Key([]string{}), rv); err != nil {
  81. return err
  82. }
  83. return enc.w.Flush()
  84. }
  85. func (enc *Encoder) safeEncode(key Key, rv reflect.Value) (err error) {
  86. defer func() {
  87. if r := recover(); r != nil {
  88. if terr, ok := r.(tomlEncodeError); ok {
  89. err = terr.error
  90. return
  91. }
  92. panic(r)
  93. }
  94. }()
  95. enc.encode(key, rv)
  96. return nil
  97. }
  98. func (enc *Encoder) encode(key Key, rv reflect.Value) {
  99. // Special case. Time needs to be in ISO8601 format.
  100. // Special case. If we can marshal the type to text, then we used that.
  101. // Basically, this prevents the encoder for handling these types as
  102. // generic structs (or whatever the underlying type of a TextMarshaler is).
  103. switch rv.Interface().(type) {
  104. case time.Time, TextMarshaler:
  105. enc.keyEqElement(key, rv)
  106. return
  107. }
  108. k := rv.Kind()
  109. switch k {
  110. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32,
  111. reflect.Int64,
  112. reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32,
  113. reflect.Uint64,
  114. reflect.Float32, reflect.Float64, reflect.String, reflect.Bool:
  115. enc.keyEqElement(key, rv)
  116. case reflect.Array, reflect.Slice:
  117. if typeEqual(tomlArrayHash, tomlTypeOfGo(rv)) {
  118. enc.eArrayOfTables(key, rv)
  119. } else {
  120. enc.keyEqElement(key, rv)
  121. }
  122. case reflect.Interface:
  123. if rv.IsNil() {
  124. return
  125. }
  126. enc.encode(key, rv.Elem())
  127. case reflect.Map:
  128. if rv.IsNil() {
  129. return
  130. }
  131. enc.eTable(key, rv)
  132. case reflect.Ptr:
  133. if rv.IsNil() {
  134. return
  135. }
  136. enc.encode(key, rv.Elem())
  137. case reflect.Struct:
  138. enc.eTable(key, rv)
  139. default:
  140. panic(e("Unsupported type for key '%s': %s", key, k))
  141. }
  142. }
  143. // eElement encodes any value that can be an array element (primitives and
  144. // arrays).
  145. func (enc *Encoder) eElement(rv reflect.Value) {
  146. switch v := rv.Interface().(type) {
  147. case time.Time:
  148. // Special case time.Time as a primitive. Has to come before
  149. // TextMarshaler below because time.Time implements
  150. // encoding.TextMarshaler, but we need to always use UTC.
  151. enc.wf(v.In(time.FixedZone("UTC", 0)).Format("2006-01-02T15:04:05Z"))
  152. return
  153. case TextMarshaler:
  154. // Special case. Use text marshaler if it's available for this value.
  155. if s, err := v.MarshalText(); err != nil {
  156. encPanic(err)
  157. } else {
  158. enc.writeQuoted(string(s))
  159. }
  160. return
  161. }
  162. switch rv.Kind() {
  163. case reflect.Bool:
  164. enc.wf(strconv.FormatBool(rv.Bool()))
  165. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32,
  166. reflect.Int64:
  167. enc.wf(strconv.FormatInt(rv.Int(), 10))
  168. case reflect.Uint, reflect.Uint8, reflect.Uint16,
  169. reflect.Uint32, reflect.Uint64:
  170. enc.wf(strconv.FormatUint(rv.Uint(), 10))
  171. case reflect.Float32:
  172. enc.wf(floatAddDecimal(strconv.FormatFloat(rv.Float(), 'f', -1, 32)))
  173. case reflect.Float64:
  174. enc.wf(floatAddDecimal(strconv.FormatFloat(rv.Float(), 'f', -1, 64)))
  175. case reflect.Array, reflect.Slice:
  176. enc.eArrayOrSliceElement(rv)
  177. case reflect.Interface:
  178. enc.eElement(rv.Elem())
  179. case reflect.String:
  180. enc.writeQuoted(rv.String())
  181. default:
  182. panic(e("Unexpected primitive type: %s", rv.Kind()))
  183. }
  184. }
  185. // By the TOML spec, all floats must have a decimal with at least one
  186. // number on either side.
  187. func floatAddDecimal(fstr string) string {
  188. if !strings.Contains(fstr, ".") {
  189. return fstr + ".0"
  190. }
  191. return fstr
  192. }
  193. func (enc *Encoder) writeQuoted(s string) {
  194. enc.wf("\"%s\"", quotedReplacer.Replace(s))
  195. }
  196. func (enc *Encoder) eArrayOrSliceElement(rv reflect.Value) {
  197. length := rv.Len()
  198. enc.wf("[")
  199. for i := 0; i < length; i++ {
  200. elem := rv.Index(i)
  201. enc.eElement(elem)
  202. if i != length-1 {
  203. enc.wf(", ")
  204. }
  205. }
  206. enc.wf("]")
  207. }
  208. func (enc *Encoder) eArrayOfTables(key Key, rv reflect.Value) {
  209. if len(key) == 0 {
  210. encPanic(errNoKey)
  211. }
  212. for i := 0; i < rv.Len(); i++ {
  213. trv := rv.Index(i)
  214. if isNil(trv) {
  215. continue
  216. }
  217. panicIfInvalidKey(key)
  218. enc.newline()
  219. enc.wf("%s[[%s]]", enc.indentStr(key), key.maybeQuotedAll())
  220. enc.newline()
  221. enc.eMapOrStruct(key, trv)
  222. }
  223. }
  224. func (enc *Encoder) eTable(key Key, rv reflect.Value) {
  225. panicIfInvalidKey(key)
  226. if len(key) == 1 {
  227. // Output an extra new line between top-level tables.
  228. // (The newline isn't written if nothing else has been written though.)
  229. enc.newline()
  230. }
  231. if len(key) > 0 {
  232. enc.wf("%s[%s]", enc.indentStr(key), key.maybeQuotedAll())
  233. enc.newline()
  234. }
  235. enc.eMapOrStruct(key, rv)
  236. }
  237. func (enc *Encoder) eMapOrStruct(key Key, rv reflect.Value) {
  238. switch rv := eindirect(rv); rv.Kind() {
  239. case reflect.Map:
  240. enc.eMap(key, rv)
  241. case reflect.Struct:
  242. enc.eStruct(key, rv)
  243. default:
  244. panic("eTable: unhandled reflect.Value Kind: " + rv.Kind().String())
  245. }
  246. }
  247. func (enc *Encoder) eMap(key Key, rv reflect.Value) {
  248. rt := rv.Type()
  249. if rt.Key().Kind() != reflect.String {
  250. encPanic(errNonString)
  251. }
  252. // Sort keys so that we have deterministic output. And write keys directly
  253. // underneath this key first, before writing sub-structs or sub-maps.
  254. var mapKeysDirect, mapKeysSub []string
  255. for _, mapKey := range rv.MapKeys() {
  256. k := mapKey.String()
  257. if typeIsHash(tomlTypeOfGo(rv.MapIndex(mapKey))) {
  258. mapKeysSub = append(mapKeysSub, k)
  259. } else {
  260. mapKeysDirect = append(mapKeysDirect, k)
  261. }
  262. }
  263. var writeMapKeys = func(mapKeys []string) {
  264. sort.Strings(mapKeys)
  265. for _, mapKey := range mapKeys {
  266. mrv := rv.MapIndex(reflect.ValueOf(mapKey))
  267. if isNil(mrv) {
  268. // Don't write anything for nil fields.
  269. continue
  270. }
  271. enc.encode(key.add(mapKey), mrv)
  272. }
  273. }
  274. writeMapKeys(mapKeysDirect)
  275. writeMapKeys(mapKeysSub)
  276. }
  277. func (enc *Encoder) eStruct(key Key, rv reflect.Value) {
  278. // Write keys for fields directly under this key first, because if we write
  279. // a field that creates a new table, then all keys under it will be in that
  280. // table (not the one we're writing here).
  281. rt := rv.Type()
  282. var fieldsDirect, fieldsSub [][]int
  283. var addFields func(rt reflect.Type, rv reflect.Value, start []int)
  284. addFields = func(rt reflect.Type, rv reflect.Value, start []int) {
  285. for i := 0; i < rt.NumField(); i++ {
  286. f := rt.Field(i)
  287. // skip unexporded fields
  288. if f.PkgPath != "" {
  289. continue
  290. }
  291. frv := rv.Field(i)
  292. if f.Anonymous {
  293. frv := eindirect(frv)
  294. t := frv.Type()
  295. if t.Kind() != reflect.Struct {
  296. encPanic(errAnonNonStruct)
  297. }
  298. addFields(t, frv, f.Index)
  299. } else if typeIsHash(tomlTypeOfGo(frv)) {
  300. fieldsSub = append(fieldsSub, append(start, f.Index...))
  301. } else {
  302. fieldsDirect = append(fieldsDirect, append(start, f.Index...))
  303. }
  304. }
  305. }
  306. addFields(rt, rv, nil)
  307. var writeFields = func(fields [][]int) {
  308. for _, fieldIndex := range fields {
  309. sft := rt.FieldByIndex(fieldIndex)
  310. sf := rv.FieldByIndex(fieldIndex)
  311. if isNil(sf) {
  312. // Don't write anything for nil fields.
  313. continue
  314. }
  315. keyName := sft.Tag.Get("toml")
  316. if keyName == "-" {
  317. continue
  318. }
  319. if keyName == "" {
  320. keyName = sft.Name
  321. }
  322. keyName, opts := getOptions(keyName)
  323. if _, ok := opts["omitempty"]; ok && isEmpty(sf) {
  324. continue
  325. } else if _, ok := opts["omitzero"]; ok && isZero(sf) {
  326. continue
  327. }
  328. enc.encode(key.add(keyName), sf)
  329. }
  330. }
  331. writeFields(fieldsDirect)
  332. writeFields(fieldsSub)
  333. }
  334. // tomlTypeName returns the TOML type name of the Go value's type. It is
  335. // used to determine whether the types of array elements are mixed (which is
  336. // forbidden). If the Go value is nil, then it is illegal for it to be an array
  337. // element, and valueIsNil is returned as true.
  338. // Returns the TOML type of a Go value. The type may be `nil`, which means
  339. // no concrete TOML type could be found.
  340. func tomlTypeOfGo(rv reflect.Value) tomlType {
  341. if isNil(rv) || !rv.IsValid() {
  342. return nil
  343. }
  344. switch rv.Kind() {
  345. case reflect.Bool:
  346. return tomlBool
  347. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32,
  348. reflect.Int64,
  349. reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32,
  350. reflect.Uint64:
  351. return tomlInteger
  352. case reflect.Float32, reflect.Float64:
  353. return tomlFloat
  354. case reflect.Array, reflect.Slice:
  355. if typeEqual(tomlHash, tomlArrayType(rv)) {
  356. return tomlArrayHash
  357. } else {
  358. return tomlArray
  359. }
  360. case reflect.Ptr, reflect.Interface:
  361. return tomlTypeOfGo(rv.Elem())
  362. case reflect.String:
  363. return tomlString
  364. case reflect.Map:
  365. return tomlHash
  366. case reflect.Struct:
  367. switch rv.Interface().(type) {
  368. case time.Time:
  369. return tomlDatetime
  370. case TextMarshaler:
  371. return tomlString
  372. default:
  373. return tomlHash
  374. }
  375. default:
  376. panic("unexpected reflect.Kind: " + rv.Kind().String())
  377. }
  378. }
  379. // tomlArrayType returns the element type of a TOML array. The type returned
  380. // may be nil if it cannot be determined (e.g., a nil slice or a zero length
  381. // slize). This function may also panic if it finds a type that cannot be
  382. // expressed in TOML (such as nil elements, heterogeneous arrays or directly
  383. // nested arrays of tables).
  384. func tomlArrayType(rv reflect.Value) tomlType {
  385. if isNil(rv) || !rv.IsValid() || rv.Len() == 0 {
  386. return nil
  387. }
  388. firstType := tomlTypeOfGo(rv.Index(0))
  389. if firstType == nil {
  390. encPanic(errArrayNilElement)
  391. }
  392. rvlen := rv.Len()
  393. for i := 1; i < rvlen; i++ {
  394. elem := rv.Index(i)
  395. switch elemType := tomlTypeOfGo(elem); {
  396. case elemType == nil:
  397. encPanic(errArrayNilElement)
  398. case !typeEqual(firstType, elemType):
  399. encPanic(errArrayMixedElementTypes)
  400. }
  401. }
  402. // If we have a nested array, then we must make sure that the nested
  403. // array contains ONLY primitives.
  404. // This checks arbitrarily nested arrays.
  405. if typeEqual(firstType, tomlArray) || typeEqual(firstType, tomlArrayHash) {
  406. nest := tomlArrayType(eindirect(rv.Index(0)))
  407. if typeEqual(nest, tomlHash) || typeEqual(nest, tomlArrayHash) {
  408. encPanic(errArrayNoTable)
  409. }
  410. }
  411. return firstType
  412. }
  413. func getOptions(keyName string) (string, map[string]struct{}) {
  414. opts := make(map[string]struct{})
  415. ss := strings.Split(keyName, ",")
  416. name := ss[0]
  417. if len(ss) > 1 {
  418. for _, opt := range ss {
  419. opts[opt] = struct{}{}
  420. }
  421. }
  422. return name, opts
  423. }
  424. func isZero(rv reflect.Value) bool {
  425. switch rv.Kind() {
  426. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  427. if rv.Int() == 0 {
  428. return true
  429. }
  430. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
  431. if rv.Uint() == 0 {
  432. return true
  433. }
  434. case reflect.Float32, reflect.Float64:
  435. if rv.Float() == 0.0 {
  436. return true
  437. }
  438. }
  439. return false
  440. }
  441. func isEmpty(rv reflect.Value) bool {
  442. switch rv.Kind() {
  443. case reflect.String:
  444. if len(strings.TrimSpace(rv.String())) == 0 {
  445. return true
  446. }
  447. case reflect.Array, reflect.Slice, reflect.Map:
  448. if rv.Len() == 0 {
  449. return true
  450. }
  451. }
  452. return false
  453. }
  454. func (enc *Encoder) newline() {
  455. if enc.hasWritten {
  456. enc.wf("\n")
  457. }
  458. }
  459. func (enc *Encoder) keyEqElement(key Key, val reflect.Value) {
  460. if len(key) == 0 {
  461. encPanic(errNoKey)
  462. }
  463. panicIfInvalidKey(key)
  464. enc.wf("%s%s = ", enc.indentStr(key), key.maybeQuoted(len(key)-1))
  465. enc.eElement(val)
  466. enc.newline()
  467. }
  468. func (enc *Encoder) wf(format string, v ...interface{}) {
  469. if _, err := fmt.Fprintf(enc.w, format, v...); err != nil {
  470. encPanic(err)
  471. }
  472. enc.hasWritten = true
  473. }
  474. func (enc *Encoder) indentStr(key Key) string {
  475. return strings.Repeat(enc.Indent, len(key)-1)
  476. }
  477. func encPanic(err error) {
  478. panic(tomlEncodeError{err})
  479. }
  480. func eindirect(v reflect.Value) reflect.Value {
  481. switch v.Kind() {
  482. case reflect.Ptr, reflect.Interface:
  483. return eindirect(v.Elem())
  484. default:
  485. return v
  486. }
  487. }
  488. func isNil(rv reflect.Value) bool {
  489. switch rv.Kind() {
  490. case reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice:
  491. return rv.IsNil()
  492. default:
  493. return false
  494. }
  495. }
  496. func panicIfInvalidKey(key Key) {
  497. for _, k := range key {
  498. if len(k) == 0 {
  499. encPanic(e("Key '%s' is not a valid table name. Key names "+
  500. "cannot be empty.", key.maybeQuotedAll()))
  501. }
  502. }
  503. }
  504. func isValidKeyName(s string) bool {
  505. return len(s) != 0
  506. }