github.com/hamba/avro@v1.8.0/writer.go (about)

     1  package avro
     2  
     3  import (
     4  	"encoding/binary"
     5  	"io"
     6  	"math"
     7  )
     8  
     9  // WriterFunc is a function used to customize the Writer.
    10  type WriterFunc func(w *Writer)
    11  
    12  // WithWriterConfig specifies the configuration to use with a writer.
    13  func WithWriterConfig(cfg API) WriterFunc {
    14  	return func(w *Writer) {
    15  		w.cfg = cfg.(*frozenConfig)
    16  	}
    17  }
    18  
    19  // Writer is an Avro specific io.Writer.
    20  type Writer struct {
    21  	cfg   *frozenConfig
    22  	out   io.Writer
    23  	buf   []byte
    24  	Error error
    25  }
    26  
    27  // NewWriter creates a new Writer.
    28  func NewWriter(out io.Writer, bufSize int, opts ...WriterFunc) *Writer {
    29  	writer := &Writer{
    30  		cfg:   DefaultConfig.(*frozenConfig),
    31  		out:   out,
    32  		buf:   make([]byte, 0, bufSize),
    33  		Error: nil,
    34  	}
    35  
    36  	for _, opt := range opts {
    37  		opt(writer)
    38  	}
    39  
    40  	return writer
    41  }
    42  
    43  // Reset resets the Writer with a new io.Writer attached.
    44  func (w *Writer) Reset(out io.Writer) {
    45  	w.out = out
    46  	w.buf = w.buf[:0]
    47  }
    48  
    49  // Buffered returns the number of buffered bytes.
    50  func (w *Writer) Buffered() int {
    51  	return len(w.buf)
    52  }
    53  
    54  // Buffer gets the Writer buffer.
    55  func (w *Writer) Buffer() []byte {
    56  	return w.buf
    57  }
    58  
    59  // Flush writes any buffered data to the underlying io.Writer.
    60  func (w *Writer) Flush() error {
    61  	if w.out == nil {
    62  		return nil
    63  	}
    64  
    65  	if w.Error != nil {
    66  		return w.Error
    67  	}
    68  
    69  	n, err := w.out.Write(w.buf)
    70  	if err != nil {
    71  		if w.Error == nil {
    72  			w.Error = err
    73  		}
    74  
    75  		return err
    76  	}
    77  
    78  	w.buf = w.buf[n:]
    79  
    80  	return nil
    81  }
    82  
    83  func (w *Writer) writeByte(b byte) {
    84  	w.buf = append(w.buf, b)
    85  }
    86  
    87  // Write writes raw bytes to the Writer.
    88  func (w *Writer) Write(b []byte) {
    89  	w.buf = append(w.buf, b...)
    90  }
    91  
    92  // WriteBool writes a Bool to the Writer.
    93  func (w *Writer) WriteBool(b bool) {
    94  	if b {
    95  		w.writeByte(0x01)
    96  		return
    97  	}
    98  
    99  	w.writeByte(0x00)
   100  }
   101  
   102  // WriteInt writes an Int to the Writer.
   103  func (w *Writer) WriteInt(i int32) {
   104  	e := uint64((uint32(i) << 1) ^ uint32(i>>31))
   105  
   106  	w.encodeInt(e)
   107  }
   108  
   109  // WriteLong writes a Long to the Writer.
   110  func (w *Writer) WriteLong(i int64) {
   111  	e := (uint64(i) << 1) ^ uint64(i>>63)
   112  
   113  	w.encodeInt(e)
   114  }
   115  
   116  func (w *Writer) encodeInt(i uint64) {
   117  	if i == 0 {
   118  		w.writeByte(0)
   119  		return
   120  	}
   121  
   122  	for i > 0 {
   123  		b := byte(i) & 0x7F
   124  		i >>= 7
   125  
   126  		if i != 0 {
   127  			b |= 0x80
   128  		}
   129  
   130  		w.writeByte(b)
   131  	}
   132  }
   133  
   134  // WriteFloat writes a Float to the Writer.
   135  func (w *Writer) WriteFloat(f float32) {
   136  	b := make([]byte, 4)
   137  	binary.LittleEndian.PutUint32(b, math.Float32bits(f))
   138  
   139  	w.buf = append(w.buf, b...)
   140  }
   141  
   142  // WriteDouble writes a Double to the Writer.
   143  func (w *Writer) WriteDouble(f float64) {
   144  	b := make([]byte, 8)
   145  	binary.LittleEndian.PutUint64(b, math.Float64bits(f))
   146  
   147  	w.buf = append(w.buf, b...)
   148  }
   149  
   150  // WriteBytes writes Bytes to the Writer.
   151  func (w *Writer) WriteBytes(b []byte) {
   152  	w.WriteLong(int64(len(b)))
   153  	w.buf = append(w.buf, b...)
   154  }
   155  
   156  // WriteString reads a String to the Writer.
   157  func (w *Writer) WriteString(s string) {
   158  	w.WriteLong(int64(len(s)))
   159  	w.buf = append(w.buf, s...)
   160  }
   161  
   162  // WriteBlockHeader writes a Block Header to the Writer.
   163  func (w *Writer) WriteBlockHeader(l, s int64) {
   164  	if s > 0 {
   165  		w.WriteLong(-l)
   166  		w.WriteLong(s)
   167  		return
   168  	}
   169  
   170  	w.WriteLong(l)
   171  }
   172  
   173  // WriteBlockCB writes a block using the callback.
   174  func (w *Writer) WriteBlockCB(callback func(w *Writer) int64) int64 {
   175  	var dummyHeader [18]byte
   176  	headerStart := len(w.buf)
   177  
   178  	// Write dummy header
   179  	w.Write(dummyHeader[:])
   180  
   181  	// Write block data
   182  	capturedAt := len(w.buf)
   183  	length := callback(w)
   184  	size := int64(len(w.buf) - capturedAt)
   185  
   186  	// Take a reference to the block data
   187  	captured := w.buf[capturedAt:len(w.buf)]
   188  
   189  	// Rewrite the header
   190  	w.buf = w.buf[:headerStart]
   191  	w.WriteBlockHeader(length, size)
   192  
   193  	// Copy the block data back to its position
   194  	w.buf = append(w.buf, captured...)
   195  
   196  	return length
   197  }