github.com/grailbio/base@v0.0.11/tsv/row_writer.go (about)

     1  // Copyright 2019 GRAIL, Inc.  All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package tsv
     6  
     7  import (
     8  	"fmt"
     9  	"io"
    10  	"reflect"
    11  	"unsafe"
    12  )
    13  
    14  // RowWriter writes structs to TSV files using field names or "tsv" tags
    15  // as TSV column headers.
    16  //
    17  // TODO: Consider letting the caller filter or reorder columns.
    18  type RowWriter struct {
    19  	w               Writer
    20  	headerDone      bool
    21  	cachedRowType   reflect.Type
    22  	cachedRowFormat rowFormat
    23  }
    24  
    25  // NewRowWriter constructs a writer.
    26  //
    27  // User must call Flush() after last Write().
    28  func NewRowWriter(w io.Writer) *RowWriter {
    29  	return &RowWriter{w: *NewWriter(w)}
    30  }
    31  
    32  // Write writes a TSV row containing the values of v's exported fields.
    33  // v must be a pointer to a struct.
    34  //
    35  // On first Write, a TSV header row is written using v's type.
    36  // Subsequent Write()s may pass v of different type, but no guarantees are made
    37  // about consistent column ordering with different types.
    38  //
    39  // By default, the column name is the struct's field name, but you can
    40  // override it by setting `tsv:"columnname"` tag in the field.
    41  //
    42  // You can optionally specify an fmt option in the tag which will control how
    43  // to format the value using the fmt package. Note that the reader may not
    44  // support all the verbs. Without the fmt option, formatting options are preset
    45  // for each type. Using the fmt option may lead to slower performance.
    46  //
    47  // Embedded structs are supported, and the default column name for nested
    48  // fields will be the unqualified name of the field.
    49  func (w *RowWriter) Write(v interface{}) error {
    50  	typ := reflect.TypeOf(v)
    51  	if typ != w.cachedRowType {
    52  		rowFormat, err := parseRowFormat(typ)
    53  		if err != nil {
    54  			return err
    55  		}
    56  		w.cachedRowType = typ
    57  		w.cachedRowFormat = rowFormat
    58  	}
    59  	if !w.headerDone {
    60  		if err := w.writeHeader(); err != nil {
    61  			return err
    62  		}
    63  		w.headerDone = true
    64  	}
    65  	return w.writeRow(v)
    66  }
    67  
    68  // Flush flushes all previously-written rows.
    69  func (w *RowWriter) Flush() error {
    70  	return w.w.Flush()
    71  }
    72  
    73  func (w *RowWriter) writeHeader() error {
    74  	for _, col := range w.cachedRowFormat {
    75  		w.w.WriteString(col.columnName)
    76  	}
    77  	return w.w.EndLine()
    78  }
    79  
    80  func (w *RowWriter) writeRow(v interface{}) error {
    81  	p := unsafe.Pointer(reflect.ValueOf(v).Pointer())
    82  	for _, col := range w.cachedRowFormat {
    83  		if col.fmt != "" {
    84  			var (
    85  				typ1 = col.typ
    86  				p1   = unsafe.Pointer(uintptr(p) + col.offset)
    87  				v    = reflect.Indirect(reflect.NewAt(typ1, p1))
    88  			)
    89  			w.w.WriteString(fmt.Sprintf("%"+col.fmt, v))
    90  			continue
    91  		}
    92  		switch col.kind {
    93  		case reflect.Bool:
    94  			v := *(*bool)(unsafe.Pointer(uintptr(p) + col.offset))
    95  			if v {
    96  				w.w.WriteString("true")
    97  			} else {
    98  				w.w.WriteString("false")
    99  			}
   100  		case reflect.String:
   101  			v := *(*string)(unsafe.Pointer(uintptr(p) + col.offset))
   102  			w.w.WriteString(v)
   103  		case reflect.Int8:
   104  			v := *(*int8)(unsafe.Pointer(uintptr(p) + col.offset))
   105  			w.w.WriteInt64(int64(v))
   106  		case reflect.Int16:
   107  			v := *(*int16)(unsafe.Pointer(uintptr(p) + col.offset))
   108  			w.w.WriteInt64(int64(v))
   109  		case reflect.Int32:
   110  			v := *(*int32)(unsafe.Pointer(uintptr(p) + col.offset))
   111  			w.w.WriteInt64(int64(v))
   112  		case reflect.Int64:
   113  			v := *(*int64)(unsafe.Pointer(uintptr(p) + col.offset))
   114  			w.w.WriteInt64(int64(v))
   115  		case reflect.Int:
   116  			v := *(*int)(unsafe.Pointer(uintptr(p) + col.offset))
   117  			w.w.WriteInt64(int64(v))
   118  		case reflect.Uint8:
   119  			v := *(*uint8)(unsafe.Pointer(uintptr(p) + col.offset))
   120  			w.w.WriteUint64(uint64(v))
   121  		case reflect.Uint16:
   122  			v := *(*uint16)(unsafe.Pointer(uintptr(p) + col.offset))
   123  			w.w.WriteUint64(uint64(v))
   124  		case reflect.Uint32:
   125  			v := *(*uint32)(unsafe.Pointer(uintptr(p) + col.offset))
   126  			w.w.WriteUint64(uint64(v))
   127  		case reflect.Uint64:
   128  			v := *(*uint64)(unsafe.Pointer(uintptr(p) + col.offset))
   129  			w.w.WriteUint64(uint64(v))
   130  		case reflect.Uint:
   131  			v := *(*uint)(unsafe.Pointer(uintptr(p) + col.offset))
   132  			w.w.WriteUint64(uint64(v))
   133  		case reflect.Float32:
   134  			v := *(*float32)(unsafe.Pointer(uintptr(p) + col.offset))
   135  			w.w.WriteFloat64(float64(v), 'g', -1)
   136  		case reflect.Float64:
   137  			v := *(*float64)(unsafe.Pointer(uintptr(p) + col.offset))
   138  			w.w.WriteFloat64(v, 'g', -1)
   139  		default:
   140  			return fmt.Errorf("unsupported type %v", col.kind)
   141  		}
   142  	}
   143  	return w.w.EndLine()
   144  }