You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

140 lines
2.6 KiB

  1. package msgpack
  2. import (
  3. "bytes"
  4. "io"
  5. "reflect"
  6. "time"
  7. "github.com/vmihailenco/msgpack/codes"
  8. )
  9. type writer interface {
  10. io.Writer
  11. WriteByte(byte) error
  12. WriteString(string) (int, error)
  13. }
  14. type byteWriter struct {
  15. io.Writer
  16. }
  17. func (w byteWriter) WriteByte(b byte) error {
  18. _, err := w.Write([]byte{b})
  19. return err
  20. }
  21. func (w byteWriter) WriteString(s string) (int, error) {
  22. return w.Write([]byte(s))
  23. }
  24. // Marshal returns the MessagePack encoding of v.
  25. func Marshal(v ...interface{}) ([]byte, error) {
  26. var buf bytes.Buffer
  27. err := NewEncoder(&buf).Encode(v...)
  28. return buf.Bytes(), err
  29. }
  30. type Encoder struct {
  31. w writer
  32. buf []byte
  33. sortMapKeys bool
  34. structAsArray bool
  35. }
  36. // NewEncoder returns a new encoder that writes to w.
  37. func NewEncoder(w io.Writer) *Encoder {
  38. bw, ok := w.(writer)
  39. if !ok {
  40. bw = byteWriter{Writer: w}
  41. }
  42. return &Encoder{
  43. w: bw,
  44. buf: make([]byte, 9),
  45. }
  46. }
  47. // SortMapKeys causes the Encoder to encode map keys in increasing order.
  48. // Supported map types are:
  49. // - map[string]string
  50. // - map[string]interface{}
  51. func (e *Encoder) SortMapKeys(v bool) *Encoder {
  52. e.sortMapKeys = v
  53. return e
  54. }
  55. // StructAsArray causes the Encoder to encode Go structs as MessagePack arrays.
  56. func (e *Encoder) StructAsArray(v bool) *Encoder {
  57. e.structAsArray = v
  58. return e
  59. }
  60. func (e *Encoder) Encode(v ...interface{}) error {
  61. for _, vv := range v {
  62. if err := e.encode(vv); err != nil {
  63. return err
  64. }
  65. }
  66. return nil
  67. }
  68. func (e *Encoder) encode(v interface{}) error {
  69. switch v := v.(type) {
  70. case nil:
  71. return e.EncodeNil()
  72. case string:
  73. return e.EncodeString(v)
  74. case []byte:
  75. return e.EncodeBytes(v)
  76. case int:
  77. return e.EncodeInt(int64(v))
  78. case int64:
  79. return e.EncodeInt(v)
  80. case uint:
  81. return e.EncodeUint(uint64(v))
  82. case uint64:
  83. return e.EncodeUint(v)
  84. case bool:
  85. return e.EncodeBool(v)
  86. case float32:
  87. return e.EncodeFloat32(v)
  88. case float64:
  89. return e.EncodeFloat64(v)
  90. case time.Duration:
  91. return e.EncodeInt(int64(v))
  92. case time.Time:
  93. return e.EncodeTime(v)
  94. }
  95. return e.EncodeValue(reflect.ValueOf(v))
  96. }
  97. func (e *Encoder) EncodeValue(v reflect.Value) error {
  98. encode := getEncoder(v.Type())
  99. return encode(e, v)
  100. }
  101. func (e *Encoder) EncodeNil() error {
  102. return e.writeCode(codes.Nil)
  103. }
  104. func (e *Encoder) EncodeBool(value bool) error {
  105. if value {
  106. return e.writeCode(codes.True)
  107. }
  108. return e.writeCode(codes.False)
  109. }
  110. func (e *Encoder) writeCode(c codes.Code) error {
  111. return e.w.WriteByte(byte(c))
  112. }
  113. func (e *Encoder) write(b []byte) error {
  114. _, err := e.w.Write(b)
  115. return err
  116. }
  117. func (e *Encoder) writeString(s string) error {
  118. _, err := e.w.WriteString(s)
  119. return err
  120. }