encode.go 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398
  1. /*
  2. Licensed under the Apache License, Version 2.0 (the "License");
  3. you may not use this file except in compliance with the License.
  4. You may obtain a copy of the License at
  5. http://www.apache.org/licenses/LICENSE-2.0
  6. Unless required by applicable law or agreed to in writing, software
  7. distributed under the License is distributed on an "AS IS" BASIS,
  8. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  9. See the License for the specific language governing permissions and
  10. limitations under the License.
  11. */
  12. package candiedyaml
  13. import (
  14. "bytes"
  15. "encoding/base64"
  16. "io"
  17. "math"
  18. "reflect"
  19. "regexp"
  20. "sort"
  21. "strconv"
  22. "time"
  23. )
  24. var (
  25. timeTimeType = reflect.TypeOf(time.Time{})
  26. marshalerType = reflect.TypeOf(new(Marshaler)).Elem()
  27. numberType = reflect.TypeOf(Number(""))
  28. nonPrintable = regexp.MustCompile("[^\t\n\r\u0020-\u007E\u0085\u00A0-\uD7FF\uE000-\uFFFD]")
  29. multiline = regexp.MustCompile("\n|\u0085|\u2028|\u2029")
  30. shortTags = map[string]string{
  31. yaml_NULL_TAG: "!!null",
  32. yaml_BOOL_TAG: "!!bool",
  33. yaml_STR_TAG: "!!str",
  34. yaml_INT_TAG: "!!int",
  35. yaml_FLOAT_TAG: "!!float",
  36. yaml_TIMESTAMP_TAG: "!!timestamp",
  37. yaml_SEQ_TAG: "!!seq",
  38. yaml_MAP_TAG: "!!map",
  39. yaml_BINARY_TAG: "!!binary",
  40. }
  41. )
  42. type Marshaler interface {
  43. MarshalYAML() (tag string, value interface{}, err error)
  44. }
  45. // An Encoder writes JSON objects to an output stream.
  46. type Encoder struct {
  47. w io.Writer
  48. emitter yaml_emitter_t
  49. event yaml_event_t
  50. flow bool
  51. err error
  52. IgnoreOmitEmpty bool
  53. }
  54. func Marshal(v interface{}) ([]byte, error) {
  55. b := bytes.Buffer{}
  56. e := NewEncoder(&b)
  57. err := e.Encode(v)
  58. return b.Bytes(), err
  59. }
  60. // NewEncoder returns a new encoder that writes to w.
  61. func NewEncoder(w io.Writer) *Encoder {
  62. e := &Encoder{w: w}
  63. yaml_emitter_initialize(&e.emitter)
  64. yaml_emitter_set_output_writer(&e.emitter, e.w)
  65. yaml_stream_start_event_initialize(&e.event, yaml_UTF8_ENCODING)
  66. e.emit()
  67. yaml_document_start_event_initialize(&e.event, nil, nil, true)
  68. e.emit()
  69. return e
  70. }
  71. func (e *Encoder) Encode(v interface{}) (err error) {
  72. defer recovery(&err)
  73. if e.err != nil {
  74. return e.err
  75. }
  76. e.marshal("", reflect.ValueOf(v), true)
  77. yaml_document_end_event_initialize(&e.event, true)
  78. e.emit()
  79. e.emitter.open_ended = false
  80. yaml_stream_end_event_initialize(&e.event)
  81. e.emit()
  82. return nil
  83. }
  84. func (e *Encoder) emit() {
  85. if !yaml_emitter_emit(&e.emitter, &e.event) {
  86. panic("bad emit")
  87. }
  88. }
  89. func (e *Encoder) marshal(tag string, v reflect.Value, allowAddr bool) {
  90. vt := v.Type()
  91. if vt.Implements(marshalerType) {
  92. e.emitMarshaler(tag, v)
  93. return
  94. }
  95. if vt.Kind() != reflect.Ptr && allowAddr {
  96. if reflect.PtrTo(vt).Implements(marshalerType) {
  97. e.emitAddrMarshaler(tag, v)
  98. return
  99. }
  100. }
  101. switch v.Kind() {
  102. case reflect.Interface:
  103. if v.IsNil() {
  104. e.emitNil()
  105. } else {
  106. e.marshal(tag, v.Elem(), allowAddr)
  107. }
  108. case reflect.Map:
  109. e.emitMap(tag, v)
  110. case reflect.Ptr:
  111. if v.IsNil() {
  112. e.emitNil()
  113. } else {
  114. e.marshal(tag, v.Elem(), true)
  115. }
  116. case reflect.Struct:
  117. e.emitStruct(tag, v)
  118. case reflect.Slice:
  119. e.emitSlice(tag, v)
  120. case reflect.String:
  121. e.emitString(tag, v)
  122. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  123. e.emitInt(tag, v)
  124. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
  125. e.emitUint(tag, v)
  126. case reflect.Float32, reflect.Float64:
  127. e.emitFloat(tag, v)
  128. case reflect.Bool:
  129. e.emitBool(tag, v)
  130. default:
  131. panic("Can't marshal type yet: " + v.Type().String())
  132. }
  133. }
  134. func (e *Encoder) emitMap(tag string, v reflect.Value) {
  135. e.mapping(tag, func() {
  136. var keys stringValues = v.MapKeys()
  137. sort.Sort(keys)
  138. for _, k := range keys {
  139. e.marshal("", k, true)
  140. e.marshal("", v.MapIndex(k), true)
  141. }
  142. })
  143. }
  144. func (e *Encoder) emitStruct(tag string, v reflect.Value) {
  145. if v.Type() == timeTimeType {
  146. e.emitTime(tag, v)
  147. return
  148. }
  149. fields := cachedTypeFields(v.Type())
  150. e.mapping(tag, func() {
  151. for _, f := range fields {
  152. fv := fieldByIndex(v, f.index)
  153. if !fv.IsValid() || f.omitEmpty && isEmptyValue(fv) && !e.IgnoreOmitEmpty {
  154. continue
  155. }
  156. e.marshal("", reflect.ValueOf(f.name), true)
  157. e.flow = f.flow
  158. e.marshal("", fv, true)
  159. }
  160. })
  161. }
  162. func (e *Encoder) emitTime(tag string, v reflect.Value) {
  163. t := v.Interface().(time.Time)
  164. bytes, _ := t.MarshalText()
  165. e.emitScalar(string(bytes), "", tag, yaml_PLAIN_SCALAR_STYLE)
  166. }
  167. func isEmptyValue(v reflect.Value) bool {
  168. switch v.Kind() {
  169. case reflect.Array, reflect.Map, reflect.Slice, reflect.String:
  170. return v.Len() == 0
  171. case reflect.Bool:
  172. return !v.Bool()
  173. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  174. return v.Int() == 0
  175. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
  176. return v.Uint() == 0
  177. case reflect.Float32, reflect.Float64:
  178. return v.Float() == 0
  179. case reflect.Interface, reflect.Ptr:
  180. return v.IsNil()
  181. }
  182. return false
  183. }
  184. func (e *Encoder) mapping(tag string, f func()) {
  185. implicit := tag == ""
  186. style := yaml_BLOCK_MAPPING_STYLE
  187. if e.flow {
  188. e.flow = false
  189. style = yaml_FLOW_MAPPING_STYLE
  190. }
  191. yaml_mapping_start_event_initialize(&e.event, nil, []byte(tag), implicit, style)
  192. e.emit()
  193. f()
  194. yaml_mapping_end_event_initialize(&e.event)
  195. e.emit()
  196. }
  197. func (e *Encoder) emitSlice(tag string, v reflect.Value) {
  198. if v.Type() == byteSliceType {
  199. e.emitBase64(tag, v)
  200. return
  201. }
  202. implicit := tag == ""
  203. style := yaml_BLOCK_SEQUENCE_STYLE
  204. if e.flow {
  205. e.flow = false
  206. style = yaml_FLOW_SEQUENCE_STYLE
  207. }
  208. yaml_sequence_start_event_initialize(&e.event, nil, []byte(tag), implicit, style)
  209. e.emit()
  210. n := v.Len()
  211. for i := 0; i < n; i++ {
  212. e.marshal("", v.Index(i), true)
  213. }
  214. yaml_sequence_end_event_initialize(&e.event)
  215. e.emit()
  216. }
  217. func (e *Encoder) emitBase64(tag string, v reflect.Value) {
  218. if v.IsNil() {
  219. e.emitNil()
  220. return
  221. }
  222. s := v.Bytes()
  223. dst := make([]byte, base64.StdEncoding.EncodedLen(len(s)))
  224. base64.StdEncoding.Encode(dst, s)
  225. e.emitScalar(string(dst), "", yaml_BINARY_TAG, yaml_DOUBLE_QUOTED_SCALAR_STYLE)
  226. }
  227. func (e *Encoder) emitString(tag string, v reflect.Value) {
  228. var style yaml_scalar_style_t
  229. s := v.String()
  230. if nonPrintable.MatchString(s) {
  231. e.emitBase64(tag, v)
  232. return
  233. }
  234. if v.Type() == numberType {
  235. style = yaml_PLAIN_SCALAR_STYLE
  236. } else {
  237. event := yaml_event_t{
  238. implicit: true,
  239. value: []byte(s),
  240. }
  241. rtag, _ := resolveInterface(event, false)
  242. if tag == "" && rtag != yaml_STR_TAG {
  243. style = yaml_DOUBLE_QUOTED_SCALAR_STYLE
  244. } else if multiline.MatchString(s) {
  245. style = yaml_LITERAL_SCALAR_STYLE
  246. } else {
  247. style = yaml_PLAIN_SCALAR_STYLE
  248. }
  249. }
  250. e.emitScalar(s, "", tag, style)
  251. }
  252. func (e *Encoder) emitBool(tag string, v reflect.Value) {
  253. s := strconv.FormatBool(v.Bool())
  254. e.emitScalar(s, "", tag, yaml_PLAIN_SCALAR_STYLE)
  255. }
  256. func (e *Encoder) emitInt(tag string, v reflect.Value) {
  257. s := strconv.FormatInt(v.Int(), 10)
  258. e.emitScalar(s, "", tag, yaml_PLAIN_SCALAR_STYLE)
  259. }
  260. func (e *Encoder) emitUint(tag string, v reflect.Value) {
  261. s := strconv.FormatUint(v.Uint(), 10)
  262. e.emitScalar(s, "", tag, yaml_PLAIN_SCALAR_STYLE)
  263. }
  264. func (e *Encoder) emitFloat(tag string, v reflect.Value) {
  265. f := v.Float()
  266. var s string
  267. switch {
  268. case math.IsNaN(f):
  269. s = ".nan"
  270. case math.IsInf(f, 1):
  271. s = "+.inf"
  272. case math.IsInf(f, -1):
  273. s = "-.inf"
  274. default:
  275. s = strconv.FormatFloat(f, 'g', -1, v.Type().Bits())
  276. }
  277. e.emitScalar(s, "", tag, yaml_PLAIN_SCALAR_STYLE)
  278. }
  279. func (e *Encoder) emitNil() {
  280. e.emitScalar("null", "", "", yaml_PLAIN_SCALAR_STYLE)
  281. }
  282. func (e *Encoder) emitScalar(value, anchor, tag string, style yaml_scalar_style_t) {
  283. implicit := tag == ""
  284. if !implicit {
  285. style = yaml_PLAIN_SCALAR_STYLE
  286. }
  287. stag := shortTags[tag]
  288. if stag == "" {
  289. stag = tag
  290. }
  291. yaml_scalar_event_initialize(&e.event, []byte(anchor), []byte(stag), []byte(value), implicit, implicit, style)
  292. e.emit()
  293. }
  294. func (e *Encoder) emitMarshaler(tag string, v reflect.Value) {
  295. if v.Kind() == reflect.Ptr && v.IsNil() {
  296. e.emitNil()
  297. return
  298. }
  299. m := v.Interface().(Marshaler)
  300. if m == nil {
  301. e.emitNil()
  302. return
  303. }
  304. t, val, err := m.MarshalYAML()
  305. if err != nil {
  306. panic(err)
  307. }
  308. if val == nil {
  309. e.emitNil()
  310. return
  311. }
  312. e.marshal(t, reflect.ValueOf(val), false)
  313. }
  314. func (e *Encoder) emitAddrMarshaler(tag string, v reflect.Value) {
  315. if !v.CanAddr() {
  316. e.marshal(tag, v, false)
  317. return
  318. }
  319. va := v.Addr()
  320. if va.IsNil() {
  321. e.emitNil()
  322. return
  323. }
  324. m := v.Interface().(Marshaler)
  325. t, val, err := m.MarshalYAML()
  326. if err != nil {
  327. panic(err)
  328. }
  329. if val == nil {
  330. e.emitNil()
  331. return
  332. }
  333. e.marshal(t, reflect.ValueOf(val), false)
  334. }