writer.go

Documentation: github.com/rs/zerolog

     1  package zerolog
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"path"
     7  	"runtime"
     8  	"strconv"
     9  	"strings"
    10  	"sync"
    11  )
    12  
    13  // LevelWriter defines as interface a writer may implement in order
    14  // to receive level information with payload.
    15  type LevelWriter interface {
    16  	io.Writer
    17  	WriteLevel(level Level, p []byte) (n int, err error)
    18  }
    19  
    20  type levelWriterAdapter struct {
    21  	io.Writer
    22  }
    23  
    24  func (lw levelWriterAdapter) WriteLevel(l Level, p []byte) (n int, err error) {
    25  	return lw.Write(p)
    26  }
    27  
    28  type syncWriter struct {
    29  	mu sync.Mutex
    30  	lw LevelWriter
    31  }
    32  
    33  // SyncWriter wraps w so that each call to Write is synchronized with a mutex.
    34  // This syncer can be used to wrap the call to writer's Write method if it is
    35  // not thread safe. Note that you do not need this wrapper for os.File Write
    36  // operations on POSIX and Windows systems as they are already thread-safe.
    37  func SyncWriter(w io.Writer) io.Writer {
    38  	if lw, ok := w.(LevelWriter); ok {
    39  		return &syncWriter{lw: lw}
    40  	}
    41  	return &syncWriter{lw: levelWriterAdapter{w}}
    42  }
    43  
    44  // Write implements the io.Writer interface.
    45  func (s *syncWriter) Write(p []byte) (n int, err error) {
    46  	s.mu.Lock()
    47  	defer s.mu.Unlock()
    48  	return s.lw.Write(p)
    49  }
    50  
    51  // WriteLevel implements the LevelWriter interface.
    52  func (s *syncWriter) WriteLevel(l Level, p []byte) (n int, err error) {
    53  	s.mu.Lock()
    54  	defer s.mu.Unlock()
    55  	return s.lw.WriteLevel(l, p)
    56  }
    57  
    58  type multiLevelWriter struct {
    59  	writers []LevelWriter
    60  }
    61  
    62  func (t multiLevelWriter) Write(p []byte) (n int, err error) {
    63  	for _, w := range t.writers {
    64  		if _n, _err := w.Write(p); err == nil {
    65  			n = _n
    66  			if _err != nil {
    67  				err = _err
    68  			} else if _n != len(p) {
    69  				err = io.ErrShortWrite
    70  			}
    71  		}
    72  	}
    73  	return n, err
    74  }
    75  
    76  func (t multiLevelWriter) WriteLevel(l Level, p []byte) (n int, err error) {
    77  	for _, w := range t.writers {
    78  		if _n, _err := w.WriteLevel(l, p); err == nil {
    79  			n = _n
    80  			if _err != nil {
    81  				err = _err
    82  			} else if _n != len(p) {
    83  				err = io.ErrShortWrite
    84  			}
    85  		}
    86  	}
    87  	return n, err
    88  }
    89  
    90  // MultiLevelWriter creates a writer that duplicates its writes to all the
    91  // provided writers, similar to the Unix tee(1) command. If some writers
    92  // implement LevelWriter, their WriteLevel method will be used instead of Write.
    93  func MultiLevelWriter(writers ...io.Writer) LevelWriter {
    94  	lwriters := make([]LevelWriter, 0, len(writers))
    95  	for _, w := range writers {
    96  		if lw, ok := w.(LevelWriter); ok {
    97  			lwriters = append(lwriters, lw)
    98  		} else {
    99  			lwriters = append(lwriters, levelWriterAdapter{w})
   100  		}
   101  	}
   102  	return multiLevelWriter{lwriters}
   103  }
   104  
   105  // TestingLog is the logging interface of testing.TB.
   106  type TestingLog interface {
   107  	Log(args ...interface{})
   108  	Logf(format string, args ...interface{})
   109  	Helper()
   110  }
   111  
   112  // TestWriter is a writer that writes to testing.TB.
   113  type TestWriter struct {
   114  	T TestingLog
   115  
   116  	// Frame skips caller frames to capture the original file and line numbers.
   117  	Frame int
   118  }
   119  
   120  // NewTestWriter creates a writer that logs to the testing.TB.
   121  func NewTestWriter(t TestingLog) TestWriter {
   122  	return TestWriter{T: t}
   123  }
   124  
   125  // Write to testing.TB.
   126  func (t TestWriter) Write(p []byte) (n int, err error) {
   127  	t.T.Helper()
   128  
   129  	n = len(p)
   130  
   131  	// Strip trailing newline because t.Log always adds one.
   132  	p = bytes.TrimRight(p, "\n")
   133  
   134  	// Try to correct the log file and line number to the caller.
   135  	if t.Frame > 0 {
   136  		_, origFile, origLine, _ := runtime.Caller(1)
   137  		_, frameFile, frameLine, ok := runtime.Caller(1 + t.Frame)
   138  		if ok {
   139  			erase := strings.Repeat("\b", len(path.Base(origFile))+len(strconv.Itoa(origLine))+3)
   140  			t.T.Logf("%s%s:%d: %s", erase, path.Base(frameFile), frameLine, p)
   141  			return n, err
   142  		}
   143  	}
   144  	t.T.Log(string(p))
   145  
   146  	return n, err
   147  }
   148  
   149  // ConsoleTestWriter creates an option that correctly sets the file frame depth for testing.TB log.
   150  func ConsoleTestWriter(t TestingLog) func(w *ConsoleWriter) {
   151  	return func(w *ConsoleWriter) {
   152  		w.Out = TestWriter{T: t, Frame: 6}
   153  	}
   154  }
   155
View as plain text