encode.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. package gocsv
  2. import (
  3. "fmt"
  4. "io"
  5. "reflect"
  6. )
  7. type encoder struct {
  8. out io.Writer
  9. }
  10. func newEncoder(out io.Writer) *encoder {
  11. return &encoder{out}
  12. }
  13. func writeFromChan(writer *SafeCSVWriter, c <-chan interface{}) error {
  14. // Get the first value. It wil determine the header structure.
  15. firstValue, ok := <-c
  16. if !ok {
  17. return fmt.Errorf("channel is closed")
  18. }
  19. inValue, inType := getConcreteReflectValueAndType(firstValue) // Get the concrete type
  20. if err := ensureStructOrPtr(inType); err != nil {
  21. return err
  22. }
  23. inInnerWasPointer := inType.Kind() == reflect.Ptr
  24. inInnerStructInfo := getStructInfo(inType) // Get the inner struct info to get CSV annotations
  25. csvHeadersLabels := make([]string, len(inInnerStructInfo.Fields))
  26. for i, fieldInfo := range inInnerStructInfo.Fields { // Used to write the header (first line) in CSV
  27. csvHeadersLabels[i] = fieldInfo.getFirstKey()
  28. }
  29. if err := writer.Write(csvHeadersLabels); err != nil {
  30. return err
  31. }
  32. write := func(val reflect.Value) error {
  33. for j, fieldInfo := range inInnerStructInfo.Fields {
  34. csvHeadersLabels[j] = ""
  35. inInnerFieldValue, err := getInnerField(val, inInnerWasPointer, fieldInfo.IndexChain) // Get the correct field header <-> position
  36. if err != nil {
  37. return err
  38. }
  39. csvHeadersLabels[j] = inInnerFieldValue
  40. }
  41. if err := writer.Write(csvHeadersLabels); err != nil {
  42. return err
  43. }
  44. return nil
  45. }
  46. if err := write(inValue); err != nil {
  47. return err
  48. }
  49. for v := range c {
  50. val, _ := getConcreteReflectValueAndType(v) // Get the concrete type (not pointer) (Slice<?> or Array<?>)
  51. if err := ensureStructOrPtr(inType); err != nil {
  52. return err
  53. }
  54. if err := write(val); err != nil {
  55. return err
  56. }
  57. }
  58. writer.Flush()
  59. return writer.Error()
  60. }
  61. func writeTo(writer *SafeCSVWriter, in interface{}, omitHeaders bool) error {
  62. inValue, inType := getConcreteReflectValueAndType(in) // Get the concrete type (not pointer) (Slice<?> or Array<?>)
  63. if err := ensureInType(inType); err != nil {
  64. return err
  65. }
  66. inInnerWasPointer, inInnerType := getConcreteContainerInnerType(inType) // Get the concrete inner type (not pointer) (Container<"?">)
  67. if err := ensureInInnerType(inInnerType); err != nil {
  68. return err
  69. }
  70. inInnerStructInfo := getStructInfo(inInnerType) // Get the inner struct info to get CSV annotations
  71. csvHeadersLabels := make([]string, len(inInnerStructInfo.Fields))
  72. for i, fieldInfo := range inInnerStructInfo.Fields { // Used to write the header (first line) in CSV
  73. csvHeadersLabels[i] = fieldInfo.getFirstKey()
  74. }
  75. if !omitHeaders {
  76. if err := writer.Write(csvHeadersLabels); err != nil {
  77. return err
  78. }
  79. }
  80. inLen := inValue.Len()
  81. for i := 0; i < inLen; i++ { // Iterate over container rows
  82. for j, fieldInfo := range inInnerStructInfo.Fields {
  83. csvHeadersLabels[j] = ""
  84. inInnerFieldValue, err := getInnerField(inValue.Index(i), inInnerWasPointer, fieldInfo.IndexChain) // Get the correct field header <-> position
  85. if err != nil {
  86. return err
  87. }
  88. csvHeadersLabels[j] = inInnerFieldValue
  89. }
  90. if err := writer.Write(csvHeadersLabels); err != nil {
  91. return err
  92. }
  93. }
  94. writer.Flush()
  95. return writer.Error()
  96. }
  97. func ensureStructOrPtr(t reflect.Type) error {
  98. switch t.Kind() {
  99. case reflect.Struct:
  100. fallthrough
  101. case reflect.Ptr:
  102. return nil
  103. }
  104. return fmt.Errorf("cannot use " + t.String() + ", only slice or array supported")
  105. }
  106. // Check if the inType is an array or a slice
  107. func ensureInType(outType reflect.Type) error {
  108. switch outType.Kind() {
  109. case reflect.Slice:
  110. fallthrough
  111. case reflect.Array:
  112. return nil
  113. }
  114. return fmt.Errorf("cannot use " + outType.String() + ", only slice or array supported")
  115. }
  116. // Check if the inInnerType is of type struct
  117. func ensureInInnerType(outInnerType reflect.Type) error {
  118. switch outInnerType.Kind() {
  119. case reflect.Struct:
  120. return nil
  121. }
  122. return fmt.Errorf("cannot use " + outInnerType.String() + ", only struct supported")
  123. }
  124. func getInnerField(outInner reflect.Value, outInnerWasPointer bool, index []int) (string, error) {
  125. oi := outInner
  126. if outInnerWasPointer {
  127. oi = outInner.Elem()
  128. }
  129. return getFieldAsString(oi.FieldByIndex(index))
  130. }