fix #85: add --auto-headers and --custom-headers (#86)

lib/parser.go

@@ -66,6 +66,43 @@ func Parse(conf cfg.Config, input io.Reader) (Tabdata, error) {
 	return data, err
 }
 
+/*
+ * Setup headers,  given headers might  be usable headers or  just the
+ * first row, which we use to  determine how many headers to generate,
+ * if enabled.
+ */
+func SetHeaders(conf cfg.Config, headers []string) []string {
+	if !conf.AutoHeaders && len(conf.CustomHeaders) == 0 {
+		return headers
+	}
+
+	if conf.AutoHeaders {
+		heads := make([]string, len(headers))
+		for idx := range headers {
+			heads[idx] = fmt.Sprintf("%d", idx+1)
+		}
+
+		return heads
+	}
+
+	if len(conf.CustomHeaders) == len(headers) {
+		return conf.CustomHeaders
+	}
+
+	// use as much custom ones we have, generate the remainder
+	heads := make([]string, len(headers))
+
+	for idx := range headers {
+		if idx < len(conf.CustomHeaders) {
+			heads[idx] = conf.CustomHeaders[idx]
+		} else {
+			heads[idx] = fmt.Sprintf("%d", idx+1)
+		}
+	}
+
+	return heads
+}
+
 /*
 Parse CSV input.
 */
@@ -87,7 +124,7 @@ func parseCSV(conf cfg.Config, input io.Reader) (Tabdata, error) {
 	}
 
 	if len(records) >= 1 {
-		data.headers = records[0]
+		data.headers = SetHeaders(conf, records[0])
 		data.columns = len(records)
 
 		for _, head := range data.headers {
@@ -98,9 +135,14 @@ func parseCSV(conf cfg.Config, input io.Reader) (Tabdata, error) {
 			}
 		}
 
-		if len(records) > 1 {
-			data.entries = records[1:]
+		if len(records) >= 1 {
+			if conf.AutoHeaders || len(conf.CustomHeaders) > 0 {
+				data.entries = records
+			} else {
+				data.entries = records[1:]
+			}
 		}
+
 	}
 
 	return data, nil
@@ -128,7 +170,9 @@ func parseTabular(conf cfg.Config, input io.Reader) (Tabdata, error) {
 			data.columns = len(parts)
 
 			// process all header fields
-			for _, part := range parts {
+			firstrow := make([]string, len(parts))
+
+			for idx, part := range parts {
 				// register widest header field
 				headerlen := len(part)
 				if headerlen > data.maxwidthHeader {
@@ -136,11 +180,22 @@ func parseTabular(conf cfg.Config, input io.Reader) (Tabdata, error) {
 				}
 
 				// register fields data
-				data.headers = append(data.headers, strings.TrimSpace(part))
+				firstrow[idx] = strings.TrimSpace(part)
 
 				// done
 				hadFirst = true
 			}
+
+			data.headers = SetHeaders(conf, firstrow)
+
+			if conf.AutoHeaders || len(conf.CustomHeaders) > 0 {
+				// we do not use generated headers, consider as row
+				if matchPattern(conf, line) == conf.InvertMatch {
+					continue
+				}
+
+				data.entries = append(data.entries, firstrow)
+			}
 		} else {
 			// data processing
 			if matchPattern(conf, line) == conf.InvertMatch {

lib/parser_test.go

@@ -366,6 +366,56 @@ func TestParserSeparators(t *testing.T) {
 	}
 }
 
+func TestParserSetHeaders(t *testing.T) {
+	row := []string{"c", "b", "c", "d", "e"}
+
+	tests := []struct {
+		name   string
+		custom []string
+		expect []string
+		auto   bool
+	}{
+		{
+			name:   "default",
+			expect: row,
+		},
+		{
+			name:   "auto",
+			expect: strings.Split("1 2 3 4 5", " "),
+			auto:   true,
+		},
+		{
+			name:   "custom-complete",
+			custom: strings.Split("A B C D E", " "),
+			expect: strings.Split("A B C D E", " "),
+		},
+		{
+			name:   "custom-too-short",
+			custom: strings.Split("A B", " "),
+			expect: strings.Split("A B 3 4 5", " "),
+		},
+		{
+			name:   "custom-too-long",
+			custom: strings.Split("A B C D E F G", " "),
+			expect: strings.Split("A B C D E", " "),
+		},
+	}
+
+	for _, testdata := range tests {
+		testname := fmt.Sprintf("parse-%s", testdata.name)
+		t.Run(testname, func(t *testing.T) {
+			conf := cfg.Config{
+				AutoHeaders:   testdata.auto,
+				CustomHeaders: testdata.custom,
+			}
+			headers := SetHeaders(conf, row)
+
+			assert.NotNil(t, headers)
+			assert.EqualValues(t, testdata.expect, headers)
+		})
+	}
+}
+
 func wrapValidateParser(conf cfg.Config, input io.Reader) (Tabdata, error) {
 	data, err := Parse(conf, input)