niki/vendor/github.com/brianvoe/gofakeit/v6/csv.go

186 lines
4.7 KiB
Go
Raw Permalink Normal View History

2024-04-26 19:30:35 +00:00
package gofakeit
import (
"bytes"
"encoding/csv"
"encoding/json"
"errors"
"fmt"
"math/rand"
"reflect"
"strings"
)
// CSVOptions defines values needed for csv generation
type CSVOptions struct {
Delimiter string `json:"delimiter" xml:"delimiter" fake:"{randomstring:[,,tab]}"`
RowCount int `json:"row_count" xml:"row_count" fake:"{number:1,10}"`
Fields []Field `json:"fields" xml:"fields" fake:"{fields}"`
}
// CSV generates an object or an array of objects in json format
// A nil CSVOptions returns a randomly structured CSV.
func CSV(co *CSVOptions) ([]byte, error) { return csvFunc(globalFaker, co) }
// CSV generates an object or an array of objects in json format
// A nil CSVOptions returns a randomly structured CSV.
func (f *Faker) CSV(co *CSVOptions) ([]byte, error) { return csvFunc(f, co) }
func csvFunc(f *Faker, co *CSVOptions) ([]byte, error) {
if co == nil {
// We didn't get a CSVOptions, so create a new random one
err := f.Struct(&co)
if err != nil {
return nil, err
}
}
// Check delimiter
if co.Delimiter == "" {
co.Delimiter = ","
}
if strings.ToLower(co.Delimiter) == "tab" {
co.Delimiter = "\t"
}
if co.Delimiter != "," && co.Delimiter != "\t" && co.Delimiter != ";" {
return nil, errors.New("invalid delimiter type")
}
// Check fields
if co.Fields == nil || len(co.Fields) <= 0 {
return nil, errors.New("must pass fields in order to build json object(s)")
}
// Make sure you set a row count
if co.RowCount <= 0 {
return nil, errors.New("must have row count")
}
b := &bytes.Buffer{}
w := csv.NewWriter(b)
w.Comma = []rune(co.Delimiter)[0]
// Add header row
header := make([]string, len(co.Fields))
for i, field := range co.Fields {
header[i] = field.Name
}
w.Write(header)
// Loop through row count +1(for header) and add fields
for i := 1; i < co.RowCount+1; i++ {
vr := make([]string, len(co.Fields))
// Loop through fields and add to them to map[string]any
for ii, field := range co.Fields {
if field.Function == "autoincrement" {
vr[ii] = fmt.Sprintf("%d", i)
continue
}
// Get function info
funcInfo := GetFuncLookup(field.Function)
if funcInfo == nil {
return nil, errors.New("invalid function, " + field.Function + " does not exist")
}
value, err := funcInfo.Generate(f.Rand, &field.Params, funcInfo)
if err != nil {
return nil, err
}
if _, ok := value.([]byte); ok {
// If it's a slice of bytes or struct, unmarshal it into an interface
var v any
if err := json.Unmarshal(value.([]byte), &v); err != nil {
return nil, err
}
value = v
}
// If the value is a list of possible values, marsha it into a string
if reflect.TypeOf(value).Kind() == reflect.Struct ||
reflect.TypeOf(value).Kind() == reflect.Ptr ||
reflect.TypeOf(value).Kind() == reflect.Map ||
reflect.TypeOf(value).Kind() == reflect.Slice {
b, err := json.Marshal(value)
if err != nil {
return nil, err
}
value = string(b)
}
vr[ii] = fmt.Sprintf("%v", value)
}
w.Write(vr)
}
w.Flush()
if err := w.Error(); err != nil {
return nil, err
}
return b.Bytes(), nil
}
func addFileCSVLookup() {
AddFuncLookup("csv", Info{
Display: "CSV",
Category: "file",
Description: "Individual lines or data entries within a CSV (Comma-Separated Values) format",
Example: `id,first_name,last_name,password
1,Markus,Moen,Dc0VYXjkWABx
2,Osborne,Hilll,XPJ9OVNbs5lm`,
Output: "[]byte",
ContentType: "text/csv",
Params: []Param{
{Field: "delimiter", Display: "Delimiter", Type: "string", Default: ",", Description: "Separator in between row values"},
{Field: "rowcount", Display: "Row Count", Type: "int", Default: "100", Description: "Number of rows"},
{Field: "fields", Display: "Fields", Type: "[]Field", Description: "Fields containing key name and function"},
},
Generate: func(r *rand.Rand, m *MapParams, info *Info) (any, error) {
co := CSVOptions{}
delimiter, err := info.GetString(m, "delimiter")
if err != nil {
return nil, err
}
co.Delimiter = delimiter
rowcount, err := info.GetInt(m, "rowcount")
if err != nil {
return nil, err
}
co.RowCount = rowcount
fieldsStr, err := info.GetStringArray(m, "fields")
if err != nil {
return nil, err
}
// Check to make sure fields has length
if len(fieldsStr) > 0 {
co.Fields = make([]Field, len(fieldsStr))
for i, f := range fieldsStr {
// Unmarshal fields string into fields array
err = json.Unmarshal([]byte(f), &co.Fields[i])
if err != nil {
return nil, err
}
}
}
f := &Faker{Rand: r}
csvOut, err := csvFunc(f, &co)
if err != nil {
return nil, err
}
return csvOut, nil
},
})
}