Merge pull request #15 from onozaty/develop/v1.9.0

Develop v1.9.0
onozaty · Jul 13, 2021 · 2ec9a47 · 2ec9a47
2 parents a0ada64 + 41701ce
commit 2ec9a47
Show file tree

Hide file tree

Showing 7 changed files with 931 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -8,6 +8,7 @@
 
 `csvt` consists of multiple subcommands.
 
+* [add](#add) Add column.
 * [choose](#choose) Choose columns.
 * [concat](#concat) Concat CSV files.
 * [count](#count) Count the number of records.
@@ -43,6 +44,92 @@ For example, when dealing with TSV files, change the delimiter to a tab as shown
 $ csvt count -i INPUT --delim "\t"
 ```
 
+## add
+
+Create a new CSV file by adding column to the input CSV file.
+
+The following values can be set for the new column.
+
+* Fixed value.
+* Same value as another column.
+* Value by template. As a template engine, [text/template](https://pkg.go.dev/text/template) will be used.
+
+### Usage
+
+```
+csvt add -i INPUT -c ADD_COLUMN [--value VALUE | --template TEMPLATE | --copy-column FROM_COLUMN] -o OUTPUT
+```
+
+```
+Usage:
+  csvt add [flags]
+
+Flags:
+  -i, --input string         Input CSV file path.
+  -c, --column string        Name of the column to add.
+      --value string         (optional) Fixed value to set for the added column.
+      --template string      (optional) Template for the value to be set for the added column.
+      --copy-column string   (optional) Name of the column from which the value is copied.
+  -o, --output string        Output CSV file path.
+  -h, --help                 help for add
+```
+
+### Example
+
+The contents of `input.csv`.
+
+```
+col1,col2
+1,a
+2,b
+3,c
+```
+
+Add "col3" as a new column. Set "x" as a fixed value.
+
+```
+$ csvt add -i input.csv -c col3 --value x -o output.csv
+```
+
+The contents of the created `output.csv`.
+
+```
+col1,col2,col3
+1,a,x
+2,b,x
+3,c,x
+```
+
+Add "col1x" by copying "col1".
+
+```
+$ csvt add -i input.csv -c col1x --copy-column col1 -o output.csv
+```
+
+```
+col1,col2,col1x
+1,a,1
+2,b,2
+3,c,3
+```
+
+Use the template to add a column that combines the values of "col1" and "col2".
+
+```
+$ csvt add -i input.csv -c col3 --template "{{.col1}}-{{.col2}}" -o output.csv
+```
+
+```
+col1,col2,col3
+1,a,1-a
+2,b,2-b
+3,c,3-c
+```
+
+Please refer to the following for template syntax.
+
+* https://pkg.go.dev/text/template
+
 ## choose
 
 Create a new CSV file by choosing columns from the input CSV file.
@@ -363,7 +450,7 @@ UserID,Name,Age,CompanyID
 
 Please refer to the following for the syntax of regular expressions.
 
-* https://golang.org/pkg/regexp/syntax/
+* https://pkg.go.dev/regexp/syntax
 
 ## header
 

diff --git a/cmd/add.go b/cmd/add.go
@@ -0,0 +1,186 @@
+package cmd
+
+import (
+	"fmt"
+	"io"
+	"strings"
+	"text/template"
+
+	"github.com/onozaty/csvt/csv"
+	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
+)
+
+func newAddCmd() *cobra.Command {
+
+	addCmd := &cobra.Command{
+		Use:   "add",
+		Short: "Add column",
+		RunE: func(cmd *cobra.Command, args []string) error {
+
+			format, err := getFlagBaseCsvFormat(cmd.Flags())
+			if err != nil {
+				return err
+			}
+
+			inputPath, _ := cmd.Flags().GetString("input")
+			addColumnName, _ := cmd.Flags().GetString("column")
+			staticValue, err := getFlagEscapedString(cmd.Flags(), "value") // バックスラッシュ記法を使いたい項目
+			if err != nil {
+				return err
+			}
+			templString, err := getFlagEscapedString(cmd.Flags(), "template") // バックスラッシュ記法を使いたい項目
+			if err != nil {
+				return err
+			}
+			copyColumnName, _ := cmd.Flags().GetString("copy-column")
+			outputPath, _ := cmd.Flags().GetString("output")
+
+			specifyValueCount := 0
+			if staticValue != "" {
+				specifyValueCount++
+			}
+			if templString != "" {
+				specifyValueCount++
+			}
+			if copyColumnName != "" {
+				specifyValueCount++
+			}
+			if specifyValueCount >= 2 {
+				return fmt.Errorf("not allowed to specify both --value and --template and --copy-column")
+			}
+
+			var templ *template.Template = nil
+
+			if templString != "" {
+				templ, err = template.New("template").Parse(templString)
+				if err != nil {
+					return errors.WithMessage(err, "--template is invalid")
+				}
+			}
+
+			// 引数の解析に成功した時点で、エラーが起きてもUsageは表示しない
+			cmd.SilenceUsage = true
+
+			return runAdd(
+				format,
+				inputPath,
+				addColumnName,
+				outputPath,
+				AddOptions{
+					staticValue:    staticValue,
+					template:       templ,
+					copyColumnName: copyColumnName,
+				})
+		},
+	}
+
+	addCmd.Flags().StringP("input", "i", "", "Input CSV file path.")
+	addCmd.MarkFlagRequired("input")
+	addCmd.Flags().StringP("column", "c", "", "Name of the column to add.")
+	addCmd.MarkFlagRequired("column")
+	addCmd.Flags().StringP("value", "", "", "(optional) Fixed value to set for the added column.")
+	addCmd.Flags().StringP("template", "", "", "(optional) Template for the value to be set for the added column.")
+	addCmd.Flags().StringP("copy-column", "", "", "(optional) Name of the column from which the value is copied.")
+	addCmd.Flags().StringP("output", "o", "", "Output CSV file path.")
+	addCmd.MarkFlagRequired("output")
+
+	return addCmd
+}
+
+type AddOptions struct {
+	staticValue    string
+	template       *template.Template
+	copyColumnName string
+}
+
+func runAdd(format csv.Format, inputPath string, addColumnName string, outputPath string, options AddOptions) error {
+
+	reader, writer, close, err := setupInputOutput(inputPath, outputPath, format)
+	if err != nil {
+		return err
+	}
+	defer close()
+
+	err = add(reader, addColumnName, writer, options)
+
+	if err != nil {
+		return err
+	}
+
+	return writer.Flush()
+}
+
+func add(reader csv.CsvReader, addColumnName string, writer csv.CsvWriter, options AddOptions) error {
+
+	// ヘッダ
+	columnNames, err := reader.Read()
+	if err != nil {
+		return errors.Wrap(err, "failed to read the CSV file")
+	}
+
+	copyColumnIndex := -1
+	if options.copyColumnName != "" {
+		copyColumnIndex, err = getTargetColumnIndex(columnNames, options.copyColumnName)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = writer.Write(append(columnNames, addColumnName))
+	if err != nil {
+		return err
+	}
+
+	// ヘッダ以外
+	for {
+		row, err := reader.Read()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return errors.Wrap(err, "failed to read the CSV file")
+		}
+
+		addColumnValue := ""
+		if options.staticValue != "" {
+
+			addColumnValue = options.staticValue
+
+		} else if options.template != nil {
+
+			addColumnValue, err = executeTemplate(options.template, columnNames, row)
+			if err != nil {
+				return err
+			}
+
+		} else if options.copyColumnName != "" {
+
+			addColumnValue = row[copyColumnIndex]
+		}
+
+		err = writer.Write(append(row, addColumnValue))
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func executeTemplate(templ *template.Template, columnNames []string, row []string) (string, error) {
+
+	// テンプレートにはヘッダ名をキーとしたmapで
+	rowMap := make(map[string]string)
+	for i, columnName := range columnNames {
+		rowMap[columnName] = row[i]
+	}
+
+	w := new(strings.Builder)
+	err := templ.Execute(w, rowMap)
+	if err != nil {
+		return "", err
+	}
+
+	return w.String(), nil
+}