github.com/lmorg/murex@v0.0.0-20240217211045-e081c89cd4ef/builtins/types/csv-bad/csv.go (about)

     1  package csvbad
     2  
     3  import (
     4  	"bufio"
     5  	"io"
     6  	"strconv"
     7  	"strings"
     8  
     9  	"github.com/lmorg/murex/config"
    10  	"github.com/lmorg/murex/lang/types"
    11  )
    12  
    13  // Parser is the CSV parser settings
    14  type Parser struct {
    15  	reader    io.Reader
    16  	Separator byte
    17  	Quote     byte
    18  	Comment   byte
    19  	Headings  bool
    20  }
    21  
    22  // NewParser creates a new CSV reader and writer - albeit it doesn't conform to Go's io.Reader / io.Writer interface{}.
    23  // The sensible thing might have been to create this as a marshaller but it's written now and works so little point
    24  // breaking it at this point in time.
    25  func NewParser(reader io.Reader, config *config.Config) (parser *Parser, err error) {
    26  	parser = new(Parser)
    27  	parser.reader = reader
    28  
    29  	parser.Separator = ','
    30  	parser.Quote = '"'
    31  	parser.Comment = '#'
    32  
    33  	v, err := config.Get("csv", "separator", types.String)
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  	if len(v.(string)) > 0 {
    38  		parser.Separator = v.(string)[0]
    39  	}
    40  
    41  	v, err = config.Get("csv", "comment", types.String)
    42  	if err != nil {
    43  		return nil, err
    44  	}
    45  	if len(v.(string)) > 0 {
    46  		parser.Comment = v.(string)[0]
    47  	}
    48  
    49  	v, err = config.Get("csv", "headings", types.Boolean)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	parser.Headings = v.(bool)
    54  	return
    55  }
    56  
    57  // ReadLine - read a line from a CSV file
    58  func (parser *Parser) ReadLine(callback func(records []string, headings []string)) (err error) {
    59  	scanner := bufio.NewScanner(parser.reader)
    60  
    61  	var headings []string
    62  
    63  	for scanner.Scan() {
    64  		var (
    65  			records     []string
    66  			current     []byte
    67  			escape      bool
    68  			quoted      bool
    69  			relativePos int
    70  		)
    71  
    72  		for _, b := range scanner.Bytes() {
    73  			switch b {
    74  			case parser.Comment:
    75  				switch {
    76  				case relativePos == 0:
    77  					break
    78  				default:
    79  					current = append(current, b)
    80  					relativePos++
    81  				}
    82  
    83  			case parser.Quote:
    84  				switch {
    85  				case escape:
    86  					current = append(current, b)
    87  					escape = false
    88  					relativePos++
    89  				case quoted:
    90  					quoted = false
    91  				case relativePos == 0:
    92  					quoted = true
    93  				default:
    94  					current = append(current, b)
    95  					relativePos++
    96  				}
    97  
    98  			case '\\':
    99  				switch {
   100  				case escape:
   101  					current = append(current, b)
   102  					escape = false
   103  					relativePos++
   104  				default:
   105  					escape = true
   106  				}
   107  
   108  			case parser.Separator:
   109  				switch {
   110  				case escape, quoted:
   111  					current = append(current, b)
   112  					escape = false
   113  					relativePos++
   114  				default:
   115  					records = append(records, string(current))
   116  					current = make([]byte, 0)
   117  					relativePos = 0
   118  				}
   119  
   120  			default:
   121  				current = append(current, b)
   122  				escape = false
   123  				relativePos++
   124  			}
   125  		}
   126  
   127  		if len(current) > 0 {
   128  			records = append(records, string(current))
   129  		}
   130  
   131  		if len(headings) == 0 {
   132  			if parser.Headings {
   133  				headings = records
   134  			} else {
   135  				for i := range records {
   136  					headings = append(headings, strconv.Itoa(i))
   137  				}
   138  			}
   139  			continue
   140  		}
   141  
   142  		if len(records) > len(headings) {
   143  			for i := len(headings); i < len(records); i++ {
   144  				headings = append(headings, strconv.Itoa(i))
   145  			}
   146  		}
   147  
   148  		callback(records, headings)
   149  	}
   150  
   151  	err = scanner.Err()
   152  	return
   153  }
   154  
   155  // ArrayToCsv marshals a list into a CSV line
   156  func (parser *Parser) ArrayToCsv(array []string) (csv []byte) {
   157  	quote := string(parser.Quote)
   158  	escapedQuote := `\` + quote
   159  	separator := quote + string(parser.Separator) + quote
   160  
   161  	for i := range array {
   162  		array[i] = strings.Replace(array[i], `\`, `\\`, -1)
   163  		array[i] = strings.Replace(array[i], "\n", `\n`, -1)
   164  		array[i] = strings.Replace(array[i], "\r", `\r`, -1)
   165  		array[i] = strings.Replace(array[i], quote, escapedQuote, -1)
   166  	}
   167  	csv = []byte(quote + strings.Join(array, separator) + quote)
   168  	return
   169  }