writer.odin 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. package csv
  2. import "core:io"
  3. import "core:strings"
  4. import "core:unicode/utf8"
  5. // Writer is a data structure used for writing records using a CSV-encoding.
  6. Writer :: struct {
  7. // Field delimiter (set to ',' with writer_init)
  8. comma: rune,
  9. // if set to true, \r\n will be used as the line terminator
  10. use_crlf: bool,
  11. w: io.Writer,
  12. }
  13. // writer_init initializes a Writer that writes to w
  14. writer_init :: proc(writer: ^Writer, w: io.Writer) {
  15. writer.comma = ',';
  16. writer.w = w;
  17. }
  18. // write writes a single CSV records to w with any of the necessarily quoting.
  19. // A record is a slice of strings, where each string is a single field.
  20. //
  21. // If the underlying io.Writer requires flushing, make sure to call io.flush
  22. write :: proc(w: ^Writer, record: []string) -> io.Error {
  23. CHAR_SET :: "\n\r\"";
  24. field_needs_quoting :: proc(w: ^Writer, field: string) -> bool {
  25. switch {
  26. case field == "": // No need to quote empty strings
  27. return false;
  28. case field == `\.`: // Postgres is weird
  29. return true;
  30. case w.comma < utf8.RUNE_SELF: // ASCII optimization
  31. for i in 0..<len(field) {
  32. switch field[i] {
  33. case '\n', '\r', '"', byte(w.comma):
  34. return true;
  35. }
  36. }
  37. case:
  38. if strings.contains_rune(field, w.comma) >= 0 {
  39. return true;
  40. }
  41. if strings.contains_any(field, CHAR_SET) {
  42. return true;
  43. }
  44. }
  45. // Leading spaces need quoting
  46. r, _ := utf8.decode_rune_in_string(field);
  47. return strings.is_space(r);
  48. }
  49. if !is_valid_delim(w.comma) {
  50. return .No_Progress; // TODO(bill): Is this a good error?
  51. }
  52. for _, field_idx in record {
  53. // NOTE(bill): declared like this so that the field can be modified later if necessary
  54. field := record[field_idx];
  55. if field_idx > 0 {
  56. if _, err := io.write_rune(w.w, w.comma); err != nil {
  57. return err;
  58. }
  59. }
  60. if !field_needs_quoting(w, field) {
  61. if _, err := io.write_string(w.w, field); err != nil {
  62. return err;
  63. }
  64. continue;
  65. }
  66. if err := io.write_byte(w.w, '"'); err != nil {
  67. return err;
  68. }
  69. for len(field) > 0 {
  70. i := strings.index_any(field, CHAR_SET);
  71. if i < 0 {
  72. i = len(field);
  73. }
  74. if _, err := io.write_string(w.w, field[:i]); err != nil {
  75. return err;
  76. }
  77. field = field[i:];
  78. if len(field) > 0 {
  79. switch field[0] {
  80. case '\r':
  81. if !w.use_crlf {
  82. if err := io.write_byte(w.w, '\r'); err != nil {
  83. return err;
  84. }
  85. }
  86. case '\n':
  87. if w.use_crlf {
  88. if _, err := io.write_string(w.w, "\r\n"); err != nil {
  89. return err;
  90. }
  91. } else {
  92. if err := io.write_byte(w.w, '\n'); err != nil {
  93. return err;
  94. }
  95. }
  96. case '"':
  97. if _, err := io.write_string(w.w, `""`); err != nil {
  98. return err;
  99. }
  100. }
  101. field = field[1:];
  102. }
  103. }
  104. if err := io.write_byte(w.w, '"'); err != nil {
  105. return err;
  106. }
  107. }
  108. if w.use_crlf {
  109. _, err := io.write_string(w.w, "\r\n");
  110. return err;
  111. }
  112. return io.write_byte(w.w, '\n');
  113. }
  114. // write_all writes multiple CSV records to w using write, and then flushes (if necessary).
  115. write_all :: proc(w: ^Writer, records: [][]string) -> io.Error {
  116. for record in records {
  117. err := write(w, record);
  118. if err != nil {
  119. return err;
  120. }
  121. }
  122. return writer_flush(w);
  123. }
  124. // writer_flush flushes the underlying io.Writer.
  125. // If the underlying io.Writer does not support flush, nil is returned.
  126. writer_flush :: proc(w: ^Writer) -> io.Error {
  127. return io.flush(auto_cast w.w);
  128. }