forked from ebhomengo/niki
186 lines
4.7 KiB
Go
186 lines
4.7 KiB
Go
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
|
|
},
|
|
})
|
|
}
|