package csvutil import ( "reflect" ) const defaultBufSize = 4096 type encField struct { field encodeFunc } type encCache struct { fields []encField buf []byte index []int record []string } func newEncCache(k typeKey, funcMap map[reflect.Type]reflect.Value, funcs []reflect.Value) (_ *encCache, err error) { fields := cachedFields(k) encFields := make([]encField, len(fields)) for i, f := range fields { fn, err := encodeFn(f.baseType, true, funcMap, funcs) if err != nil { return nil, err } encFields[i] = encField{ field: f, encodeFunc: fn, } } return &encCache{ fields: encFields, buf: make([]byte, 0, defaultBufSize), index: make([]int, len(encFields)), record: make([]string, len(encFields)), }, nil } // Encoder writes structs CSV representations to the output stream. type Encoder struct { // Tag defines which key in the struct field's tag to scan for names and // options (Default: 'csv'). Tag string // If AutoHeader is true, a struct header is encoded during the first call // to Encode automatically (Default: true). AutoHeader bool w Writer c *encCache noHeader bool typeKey typeKey funcMap map[reflect.Type]reflect.Value ifaceFuncs []reflect.Value } // NewEncoder returns a new encoder that writes to w. func NewEncoder(w Writer) *Encoder { return &Encoder{ w: w, noHeader: true, AutoHeader: true, } } // Register registers a custom encoding function for a concrete type or interface. // The argument f must be of type: // func(T) ([]byte, error) // // T must be a concrete type such as Foo or *Foo, or interface that has at // least one method. // // During encoding, fields are matched by the concrete type first. If match is not // found then Encoder looks if field implements any of the registered interfaces // in order they were registered. // // Register panics if: // - f does not match the right signature // - f is an empty interface // - f was already registered // // Register is based on the encoding/json proposal: // https://github.com/golang/go/issues/5901. func (e *Encoder) Register(f interface{}) { v := reflect.ValueOf(f) typ := v.Type() if typ.Kind() != reflect.Func || typ.NumIn() != 1 || typ.NumOut() != 2 || typ.Out(0) != _bytes || typ.Out(1) != _error { panic("csvutil: func must be of type func(T) ([]byte, error)") } argType := typ.In(0) if argType.Kind() == reflect.Interface && argType.NumMethod() == 0 { panic("csvutil: func argument type must not be an empty interface") } if e.funcMap == nil { e.funcMap = make(map[reflect.Type]reflect.Value) } if _, ok := e.funcMap[argType]; ok { panic("csvutil: func " + typ.String() + " already registered") } e.funcMap[argType] = v if argType.Kind() == reflect.Interface { e.ifaceFuncs = append(e.ifaceFuncs, v) } } // Encode writes the CSV encoding of v to the output stream. The provided // argument v must be a struct, struct slice or struct array. // // Only the exported fields will be encoded. // // First call to Encode will write a header unless EncodeHeader was called first // or AutoHeader is false. Header names can be customized by using tags // ('csv' by default), otherwise original Field names are used. // // Header and fields are written in the same order as struct fields are defined. // Embedded struct's fields are treated as if they were part of the outer struct. // Fields that are embedded types and that are tagged are treated like any // other field, but they have to implement Marshaler or encoding.TextMarshaler // interfaces. // // Marshaler interface has the priority over encoding.TextMarshaler. // // Tagged fields have the priority over non tagged fields with the same name. // // Following the Go visibility rules if there are multiple fields with the same // name (tagged or not tagged) on the same level and choice between them is // ambiguous, then all these fields will be ignored. // // Nil values will be encoded as empty strings. Same will happen if 'omitempty' // tag is set, and the value is a default value like 0, false or nil interface. // // Bool types are encoded as 'true' or 'false'. // // Float types are encoded using strconv.FormatFloat with precision -1 and 'G' // format. NaN values are encoded as 'NaN' string. // // Fields of type []byte are being encoded as base64-encoded strings. // // Fields can be excluded from encoding by using '-' tag option. // // Examples of struct tags: // // // Field appears as 'myName' header in CSV encoding. // Field int `csv:"myName"` // // // Field appears as 'Field' header in CSV encoding. // Field int // // // Field appears as 'myName' header in CSV encoding and is an empty string // // if Field is 0. // Field int `csv:"myName,omitempty"` // // // Field appears as 'Field' header in CSV encoding and is an empty string // // if Field is 0. // Field int `csv:",omitempty"` // // // Encode ignores this field. // Field int `csv:"-"` // // // Encode treats this field exactly as if it was an embedded field and adds // // "my_prefix_" to each field's name. // Field Struct `csv:"my_prefix_,inline"` // // // Encode treats this field exactly as if it was an embedded field. // Field Struct `csv:",inline"` // // Fields with inline tags that have a non-empty prefix must not be cyclic // structures. Passing such values to Encode will result in an infinite loop. // // Encode doesn't flush data. The caller is responsible for calling Flush() if // the used Writer supports it. func (e *Encoder) Encode(v interface{}) error { return e.encode(reflect.ValueOf(v)) } // EncodeHeader writes the CSV header of the provided struct value to the output // stream. The provided argument v must be a struct value. // // The first Encode method call will not write header if EncodeHeader was called // before it. This method can be called in cases when a data set could be // empty, but header is desired. // // EncodeHeader is like Header function, but it works with the Encoder and writes // directly to the output stream. Look at Header documentation for the exact // header encoding rules. func (e *Encoder) EncodeHeader(v interface{}) error { typ, err := valueType(v) if err != nil { return err } return e.encodeHeader(typ) } func (e *Encoder) encode(v reflect.Value) error { val := walkValue(v) if !val.IsValid() { return &InvalidEncodeError{} } switch val.Kind() { case reflect.Struct: return e.encodeStruct(val) case reflect.Array, reflect.Slice: if walkType(val.Type().Elem()).Kind() != reflect.Struct { return &InvalidEncodeError{v.Type()} } return e.encodeArray(val) default: return &InvalidEncodeError{v.Type()} } } func (e *Encoder) encodeStruct(v reflect.Value) error { if e.AutoHeader && e.noHeader { if err := e.encodeHeader(v.Type()); err != nil { return err } } return e.marshal(v) } func (e *Encoder) encodeArray(v reflect.Value) error { l := v.Len() for i := 0; i < l; i++ { if err := e.encodeStruct(walkValue(v.Index(i))); err != nil { return err } } return nil } func (e *Encoder) encodeHeader(typ reflect.Type) error { fields, _, _, record, err := e.cache(typ) if err != nil { return err } for i, f := range fields { record[i] = f.name } if err := e.w.Write(record); err != nil { return err } e.noHeader = false return nil } func (e *Encoder) marshal(v reflect.Value) error { fields, buf, index, record, err := e.cache(v.Type()) if err != nil { return err } for i, f := range fields { v := walkIndex(v, f.index) omitempty := f.tag.omitEmpty if v.Kind() == reflect.Ptr || v.Kind() == reflect.Interface { // We should disable omitempty for pointer and interface values, // because if it's nil we will automatically encode it as an empty // string. However, the initialized pointer should not be affected, // even if it's a default value. omitempty = false } if !v.IsValid() { index[i] = 0 continue } b, err := f.encodeFunc(buf, v, omitempty) if err != nil { return err } index[i], buf = len(b)-len(buf), b } out := string(buf) for i, n := range index { record[i], out = out[:n], out[n:] } e.c.buf = buf[:0] return e.w.Write(record) } func (e *Encoder) tag() string { if e.Tag == "" { return defaultTag } return e.Tag } func (e *Encoder) cache(typ reflect.Type) ([]encField, []byte, []int, []string, error) { if k := (typeKey{e.tag(), typ}); k != e.typeKey { c, err := newEncCache(k, e.funcMap, e.ifaceFuncs) if err != nil { return nil, nil, nil, nil, err } e.c, e.typeKey = c, k } return e.c.fields, e.c.buf[:0], e.c.index, e.c.record, nil } func walkIndex(v reflect.Value, index []int) reflect.Value { for _, i := range index { v = walkPtr(v) if !v.IsValid() { return reflect.Value{} } v = v.Field(i) } return v } func walkPtr(v reflect.Value) reflect.Value { for v.Kind() == reflect.Ptr { v = v.Elem() } return v } func walkValue(v reflect.Value) reflect.Value { for { switch v.Kind() { case reflect.Ptr, reflect.Interface: v = v.Elem() default: return v } } } func walkType(typ reflect.Type) reflect.Type { for typ.Kind() == reflect.Ptr { typ = typ.Elem() } return typ }